devbox
我家自动化
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Yolo v7的最简TensorFlow实现_yolo tensorflow

作者：我家自动化 | 2024-03-23 14:10:52

赞

踩

yolo tensorflow

Yolo v7去年推出之后，取得了很好的性能。作者也公布了基于Pytorch实现的源代码。在我之前的几篇博客当中，对代码进行了深入的解析，了解了Yolo v7的技术细节和实现机制。因为我一直是用的Tensorflow，因此也想尝试把代码移植到Tensorflow上。

数据集的构建

直接运行Yolo v7源代码里面的get_coco.sh脚本下载coco数据集，脚本代码如下：

  1. #!/bin/bash
  2. # COCO 2017 dataset http://cocodataset.org
  3. # Download command: bash ./scripts/get_coco.sh
  4.  
  5. # Download/unzip labels
  6. d='./' # unzip directory
  7. url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
  8. f='coco2017labels-segments.zip' # or 'coco2017labels.zip', 68 MB
  9. echo 'Downloading' $url$f ' ...'
  10. curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
  11.  
  12. # Download/unzip images
  13. d='./coco/images' # unzip directory
  14. url=http://images.cocodataset.org/zips/
  15. f1='train2017.zip' # 19G, 118k images
  16. f2='val2017.zip'   # 1G, 5k images
  17. f3='test2017.zip'  # 7G, 41k images (optional)
  18. for f in $f1 $f2 $f3; do
  19.   echo 'Downloading' $url$f '...'
  20.   curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
  21. done
  22. wait # finish background tasks

数据下载完成之后，在images和labels目录下分别有train2017, val2017, test2017这三个子目录，对应训练/验证/测试数据。

然后我们可以基于Tensorflow来构建一个训练的数据集，需要对训练的图像进行增强，包括了包括了Mosaic拼接，随机拷贝图像，随机形变，色彩调整等，相应的图像里面的物体Label也要做相应的变换。具体的工作原理可以见我之前的博客，解读YOLO v7的代码(二)训练数据的准备-CSDN博客

这里我定义了一个Dataloader的类，负责对训练集的数据进行相应的图像增强处理，这里的处理过程和Yolov7源码的基本是一致的，只是做了一些小的修改，就是当做了Mosaic拼接之后，如果随机形变是进行缩小，那么有可能会出现物体的检测框超出图像的情况，这里我根据物体的segments数据进行了裁减，使得不会超出图像。

对于验证集的数据，我们不需要进行图像增强，只需要对图像的长边缩放到640即可，空白部分进行padding。Tensorflow的dataset的定义如下：

  1. def map_val_fn(t: tf.Tensor):
  2.     filename = str(t.numpy(), encoding='utf-8')
  3.     imgid = int(filename[20:32])
  4.     # Load image
  5.     img, (h0, w0), (h, w) = load_image(filename)
  6.     #augment_hsv(img, hgain=hsv_h, sgain=hsv_s, vgain=hsv_v)
  7.  
  8.     # Labels
  9.     label_filename = val_label_path + filename.split('/')[-1].split('.')[0] + '.txt'
  10.     labels, _ = load_labels(label_filename)
  11.     labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, 0, 0)  # normalized xywh to pixel xyxy format
  12.     labels[:, 1:5] = xyxy2xywh(labels[:, 1:5])  # convert xyxy to xywh
  13.     labels[:, 1:5] /= img_size  # normalized height 0-1
  14.     
  15.     img = img[:, :, ::-1].transpose(2,0,1)
  16.     img = img/255.
  17.     
  18.     img_hw = tf.concat([h0, w0], axis=0)
  19.     return img, labels, img_hw, imgid
  20.  
  21. dataset_val = tf.data.Dataset.list_files("coco/images/val2017/*.jpg", shuffle=False)
  22. dataset_val = dataset_val.map(
  23.     lambda x: tf.py_function(func=map_val_fn, inp=[x], Tout=[tf.float32, tf.float32, tf.int32, tf.int32]), 
  24.     num_parallel_calls=tf.data.experimental.AUTOTUNE)
  25. dataset_val = dataset_val\
  26.     .padded_batch(val_batch_size, padded_shapes=([3, img_size, img_size], [None, 5], [2], []), padding_values=(144/255., 0., 0, 0))\
  27.     .prefetch(tf.data.experimental.AUTOTUNE)

对于训练集的dataset，本来我也是打算按类似以上验证集的方式来定义，只是把map函数替换为对应的Dataloader里面的函数，具体代码可以见dataloader.py。但是我发现这种方式效率不高，在实际测试中发现，因为这个图像增强的过程比较复杂，CPU需要花费较多的事件处理，虽然Tensorflow dataset的map和prefetch提供了一个Autotune的参数可以进行并行处理的优化，但是效果不是太理想，还是出现GPU等待CPU处理完数据的情况。为此我自己写了一个并行处理的函数，利用Python multiprocessing的多进程函数，来对图像进行并行处理，当GPU在训练100个Batch的时候，CPU并行准备下100个Batch的训练数据，这样可以大幅提高性能。

具体做法是创建一个share memory给各个子进程共享，然后在训练集的图像中随机抽取一部分文件名，分配给几个子进程，每个子进程读取这些图像，进行相应的图像处理，以及对相应的图像Label文件进行处理，并把处理后的数据写入到Share memory的对应位置。最后有一个独立的子进程对Share memory的数据进行合并整理，然后就可以基于整理后的数据直接构建一个dataset了。

相关的代码如下：

  1. #对传入的图像ID进行增强处理，并把结果写入到共享内存
  2. def augment_data(imgids, datasize, memory_name, offset, q):
  3.     dataset = Dataloader(img_size, train_image_dir, train_label_dir, imgids, hyp)
  4.     traindata = dataset.generateTrainData(datasize)
  5.     traindata_obj = pickle.dumps(traindata, protocol=pickle.HIGHEST_PROTOCOL)
  6.     existing_shm = shared_memory.SharedMemory(name=memory_name)
  7.     existing_shm.buf[offset:offset+len(traindata_obj)] = traindata_obj 
  8.     q.put((offset, offset+len(traindata_obj)))
  9.     existing_shm.close()
  10.  
  11. #对图像处理子进程的结果进行合并
  12. def merge_subprocess(q, subprocess_num, memory_name):
  13.     results = []
  14.     while(True):
  15.         msg = q.get()
  16.         if msg is not None:
  17.             results.append(msg)
  18.         if len(results)>=subprocess_num:
  19.             break
  20.         else:
  21.             time.sleep(1)
  22.     existing_shm = shared_memory.SharedMemory(name=memory_name)
  23.     merge_data =  []
  24.     for result in results:
  25.         merge_data.extend(pickle.loads(existing_shm.buf[result[0]:result[1]]))
  26.     merge_data_obj = pickle.dumps(merge_data, protocol=pickle.HIGHEST_PROTOCOL)
  27.     existing_shm.buf[:len(merge_data_obj)] = merge_data_obj
  28.     existing_shm.close()
  29.     q.put(len(merge_data_obj))
  30.  
  31.  
  32. #启动多个子进程进行图像增强处理，并对结果进行汇总整理
  33. def prepare_traindata(memory_name):
  34.     sample_imgid = sample(imgid_train, sample_len)    #随机选取一部分训练集图像的文件名
  35.     subprocess_list = []
  36.     for i in range(subprocess_num):      #启动多个子进程，分别对图像和Label进行处理
  37.         subprocess_list.append(
  38.             mp.Process(
  39.                 target=augment_data, 
  40.                 args=(sample_imgid[i*imgid_num_process:(i+1)*imgid_num_process], data_size//subprocess_num, memory_name, i*shared_memory_size//subprocess_num, q, )
  41.             )
  42.         )
  43.     for p in subprocess_list:
  44.         p.start()
  45.     #启动子进程对处理结果进行汇总整理
  46.     p0 = mp.Process(target=merge_subprocess, args=(q, subprocess_num, memory_name,))
  47.     p0.start()
  48.     return p0
  49.  
  50.  
  51. image_cache = shared_memory.SharedMemory(name="dataset", create=True, size=shared_memory_size)     #创建共享内存
  52.  
  53. merge_proc = prepare_traindata("dataset")
  54.  
  55. #等待汇总子进程执行完毕，从Queue中获取数据size，并进行反序列化
  56. merge_proc.join()
  57. msg = q.get()
  58. if msg>0:
  59.     traindata = pickle.loads(image_cache.buf[:msg])
  60. else:
  61.     print("Could not load training data.")
  62.     image_cache.close()
  63.     image_cache.unlink()
  64.  
  65. image_cache.close()
  66. image_cache.unlink()
  67.  
  68. def traindata_gen():
  69.     global traindata
  70.     i = 0
  71.     while i<len(traindata):
  72.         yield traindata[i][0]/255., traindata[i][1]
  73.         i += 1
  74.  
  75. #构建dataset
  76. dataset = tf.data.Dataset.from_generator(
  77.     traindata_gen,
  78.     output_types=(tf.float32, tf.float32), 
  79.     output_shapes=((3, img_size, img_size), (None, 5)))
  80. dataset = dataset.padded_batch(batch_size, padded_shapes=([3, img_size, img_size], [None, 5]))
  81. dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

模型的定义

构建一个YOLO v7的模型，模型的结构解读可见我之前的另一篇博客解读YOLO v7的代码(一)模型结构研究_gzroy的博客-CSDN博客

定义一个yolo.py文件，里面定义了模型的自定义层和对模型进行组装。

  1. import tensorflow as tf
  2. from tensorflow import keras
  3. l=tf.keras.layers
  4. from params import *
  5.  
  6. @tf.keras.utils.register_keras_serializable()
  7. class YoloConv(keras.layers.Layer):
  8.     def __init__(self, filters, kernel_size, strides, padding='same', bias=False, activation='swish', **kwargs):
  9.         super(YoloConv, self).__init__(**kwargs)
  10.         self.activation = activation
  11.         self.filters = filters
  12.         self.kernel_size = kernel_size
  13.         self.strides = strides
  14.         self.padding = padding
  15.         self.bias = bias
  16.         self.cv = l.Conv2D(filters=self.filters, 
  17.             kernel_size=self.kernel_size, 
  18.             strides=self.strides,
  19.             padding=self.padding,
  20.             data_format='channels_first',
  21.             use_bias=self.bias,
  22.             kernel_initializer='he_normal',
  23.             kernel_regularizer=tf.keras.regularizers.l2(l=weight_decay))
  24.         self.bn = l.BatchNormalization(axis=1)
  25.         self.swish = l.Activation('swish')
  26.  
  27.     def call(self, inputs, training):
  28.         output = self.cv(inputs)
  29.         output = self.bn(output, training)
  30.         if self.activation=='swish':
  31.             output = self.swish(output)
  32.         else:
  33.             output = output
  34.         return output
  35.  
  36.     def get_config(self):
  37.         config = super(YoloConv, self).get_config()
  38.         config.update({
  39.             "activation": self.activation,
  40.             "filters": self.filters,
  41.             "kernel_size": self.kernel_size,
  42.             "strides": self.strides,
  43.             "padding": self.padding,
  44.             "bias": self.bias
  45.         })
  46.         return config
  47.  
  48. @tf.keras.utils.register_keras_serializable()
  49. class Elan(keras.layers.Layer):
  50.     def __init__(self, filters, **kwargs):
  51.         super(Elan, self).__init__(**kwargs)
  52.         self.filters = filters
  53.         self.cv1 = YoloConv(self.filters, 1, 1)
  54.         self.cv2 = YoloConv(self.filters, 1, 1)
  55.         self.cv3 = YoloConv(self.filters, 3, 1)
  56.         self.cv4 = YoloConv(self.filters, 3, 1)
  57.         self.cv5 = YoloConv(self.filters, 3, 1)
  58.         self.cv6 = YoloConv(self.filters, 3, 1)
  59.         self.cv7 = YoloConv(self.filters*4, 1, 1)
  60.         self.concat = l.Concatenate(axis=1)
  61.         
  62.     def call(self, inputs, training):
  63.         output1 = self.cv1(inputs, training)
  64.         output2 = self.cv2(inputs, training)
  65.         output3 = self.cv4(self.cv3(output2, training), training)
  66.         output4 = self.cv6(self.cv5(output3, training), training)
  67.         output = self.concat([output1, output2, output3, output4])
  68.         output = self.cv7(output, training)
  69.         return output
  70.  
  71.     def get_config(self):
  72.         config = super(Elan, self).get_config()
  73.         config.update({
  74.             "filters": self.filters
  75.         })
  76.         return config
  77.  
  78. @tf.keras.utils.register_keras_serializable()
  79. class MP(keras.layers.Layer):
  80.     def __init__(self, filters, k=2):
  81.         super(MP, self).__init__()
  82.         self.filters = filters
  83.         self.k = k
  84.         self.cv1 = YoloConv(filters, 1, 1)
  85.         self.cv2 = YoloConv(filters, 1, 1)
  86.         self.cv3 = YoloConv(filters, 3, 2)
  87.         self.pool = l.MaxPool2D(pool_size=self.k, strides=self.k, padding='same', data_format='channels_first')
  88.         self.concat = l.Concatenate(axis=1)
  89.         
  90.     def call(self, inputs, training):
  91.         output1 = self.pool(inputs)
  92.         output1 = self.cv1(output1, training)
  93.         output2 = self.cv2(inputs, training)
  94.         output2 = self.cv3(output2, training)
  95.         output = self.concat([output1, output2])
  96.         return output
  97.  
  98.     def get_config(self):
  99.         config = super(MP, self).get_config()
  100.         config.update({
  101.             "filters": self.filters,
  102.             "k": self.k
  103.         })
  104.         return config
  105.  
  106. @tf.keras.utils.register_keras_serializable()    
  107. class SPPCSPC(keras.layers.Layer):
  108.     def __init__(self, filters, e=0.5, k=(5,9,13)):
  109.         super(SPPCSPC, self).__init__()
  110.         self.filters = filters
  111.         self.e = e
  112.         self.k = k
  113.         c_ = int(2 * self.filters * self.e)
  114.         self.cv1 = YoloConv(c_, 1, 1)
  115.         self.cv2 = YoloConv(c_, 1, 1)
  116.         self.cv3 = YoloConv(c_, 3, 1)
  117.         self.cv4 = YoloConv(c_, 1, 1)
  118.         self.m = [l.MaxPool2D(pool_size=x, strides=1, padding='same', data_format='channels_first') for x in k]
  119.         self.cv5 = YoloConv(c_, 1, 1)
  120.         self.cv6 = YoloConv(c_, 3, 1)
  121.         self.cv7 = YoloConv(filters, 1, 1)
  122.         self.concat = l.Concatenate(axis=1)
  123.         
  124.     def call(self, inputs, training):
  125.         output1 = self.cv4(self.cv3(self.cv1(inputs, training), training), training)
  126.         output2 = self.concat([output1] + [m(output1) for m in self.m])
  127.         output2 = self.cv6(self.cv5(output2, training), training)
  128.         output3 = self.cv2(inputs, training)
  129.         output = self.cv7(self.concat([output2, output3]), training)
  130.         return output
  131.     
  132.     def get_config(self):
  133.         config = super(SPPCSPC, self).get_config()
  134.         config.update({
  135.             "filters": self.filters,
  136.             "k": self.k,
  137.             "e": self.e
  138.         })
  139.         return config
  140.  
  141. @tf.keras.utils.register_keras_serializable()
  142. class Elan_A(keras.layers.Layer):
  143.     def __init__(self, filters):
  144.         super(Elan_A, self).__init__()
  145.         self.filters = filters
  146.         self.cv1 = YoloConv(filters, 1, 1)
  147.         self.cv2 = YoloConv(filters, 1, 1)
  148.         self.cv3 = YoloConv(filters//2, 3, 1)
  149.         self.cv4 = YoloConv(filters//2, 3, 1)
  150.         self.cv5 = YoloConv(filters//2, 3, 1)
  151.         self.cv6 = YoloConv(filters//2, 3, 1)
  152.         self.cv7 = YoloConv(filters, 1, 1)
  153.         self.concat = l.Concatenate(axis=1)
  154.         
  155.     def call(self, inputs, training):
  156.         output1 = self.cv1(inputs, training)
  157.         output2 = self.cv2(inputs, training)
  158.         output3 = self.cv3(output2, training)
  159.         output4 = self.cv4(output3, training)
  160.         output5 = self.cv5(output4, training)
  161.         output6 = self.cv6(output5, training)
  162.         output7 = self.concat([output1, output2, output3, output4, output5, output6])
  163.         output = self.cv7(output7, training)
  164.         return output
  165.     
  166.     def get_config(self):
  167.         config = super(Elan_A, self).get_config()
  168.         config.update({
  169.             "filters": self.filters,
  170.         })
  171.         return config
  172.  
  173. @tf.keras.utils.register_keras_serializable()
  174. class RepConv(keras.layers.Layer):
  175.     def __init__(self, filters):
  176.         super(RepConv, self).__init__()
  177.         self.filters = filters
  178.         self.cv1 = YoloConv(filters, 3, 1, activation=None)
  179.         self.cv2 = YoloConv(filters, 1, 1, activation=None)
  180.         self.swish = l.Activation('swish')
  181.         
  182.     def call(self, inputs, training):
  183.         output1 = self.cv1(inputs, training)
  184.         output2 = self.cv2(inputs, training)
  185.         output = self.swish(output1+output2)
  186.         return output
  187.     
  188.     def get_config(self):
  189.         config = super(RepConv, self).get_config()
  190.         config.update({
  191.             "filters": self.filters,
  192.         })
  193.         return config
  194.  
  195. @tf.keras.utils.register_keras_serializable()
  196. class IDetect(keras.layers.Layer):
  197.     def __init__(self, shape, no, na, grids):
  198.         super(IDetect, self).__init__()
  199.         #self.a = tf.random.normal((1,shape,1,1), mean=0.0, stddev=0.02, dtype=tf.dtypes.float16)
  200.         self.a = tf.Variable(tf.random.normal((1,shape,1,1), mean=0.0, stddev=0.02, dtype=tf.dtypes.float16))
  201.         self.m = tf.Variable(tf.random.normal((1,no*na,1,1), mean=0.0, stddev=0.02, dtype=tf.dtypes.float16))
  202.         #self.a = keras.initializers.RandomNormal(mean=0., stddev=0.02)(shape=(1,shape,1,1))
  203.         #self.m = keras.initializers.RandomNormal(mean=0., stddev=0.02)(shape=(1,no*na,1,1))
  204.         self.cv = YoloConv(no*na, 1, 1, bias=True, activation=None)
  205.         self.shape = shape
  206.         self.no = no
  207.         self.na = na
  208.         self.grids = grids
  209.         self.reshape = l.Reshape([self.na, self.no, self.grids*self.grids])
  210.         #self.permute = l.Permute([1,3,4,2])
  211.         self.permute = l.Permute([1,3,2])
  212.         self.activation = l.Activation('linear', dtype='float32')
  213.     
  214.     def call(self, inputs, training):
  215.         #output = l.Add()([inputs, self.a])
  216.         output = inputs + self.a
  217.         output = self.cv(output, training)
  218.         output = self.m * output
  219.         #output = self.cv(inputs)
  220.         #output = tf.reshape(output, [-1, self.na, self.no, self.grids, self.grids])
  221.         output = self.reshape(output)
  222.         #output = tf.transpose(output, perm=[0,1,3,4,2])
  223.         output = self.permute(output)
  224.         output = self.activation(output)
  225.         return output
  226.  
  227.     def get_config(self):
  228.         config = super(IDetect, self).get_config()
  229.         config.update({
  230.             "no": self.no,
  231.             "na": self.na,
  232.             "grids": self.grids,
  233.             "shape": self.shape
  234.         })
  235.         return config
  236.  
  237. def create_model():
  238.     inputs = keras.Input(shape=(3, img_size, img_size))
  239.     x = YoloConv(32, 3, 1)(inputs)    #[32, img_size, img_size]
  240.     x = YoloConv(64, 3, 2)(x)         #[64, img_size/2, img_size/2]
  241.     x = YoloConv(64, 3, 1)(x)         #[64, img_size/2, img_size/2]
  242.     x = YoloConv(128, 3, 2)(x)        #[128, img_size/4, img_size/4]
  243.     x = Elan(64)(x)                   #11
  244.     x = MP(128)(x)                    #16
  245.     route1 = Elan(128)(x)             #24
  246.     x = MP(256)(route1)               #29
  247.     route2 = Elan(256)(x)             #37
  248.     x = MP(512)(route2)               #42
  249.     x = Elan(256)(x)                  #50
  250.     route3 = SPPCSPC(512)(x)          #51
  251.     x = YoloConv(256, 1, 1)(route3)
  252.     x = l.UpSampling2D(size=(2, 2), data_format='channels_first', interpolation='nearest')(x)
  253.     x = l.Concatenate(axis=1)([x, YoloConv(256, 1, 1)(route2)])
  254.     route4 = Elan_A(256)(x)           #63
  255.     x = YoloConv(128, 1, 1)(route4)
  256.     x = l.UpSampling2D(size=(2, 2), data_format='channels_first', interpolation='nearest')(x)
  257.     x = l.Concatenate(axis=1)([x, YoloConv(128, 1, 1)(route1)])
  258.     route5 = Elan_A(128)(x)           #75, Connect to Detector 1
  259.     x = MP(128)(route5)  
  260.     x = l.Concatenate(axis=1)([x, route4])
  261.     route6 = Elan_A(256)(x)           #88, Connect to Detector 2
  262.     x = MP(256)(route6)   
  263.     x = l.Concatenate(axis=1)([x, route3])
  264.     route7 = Elan_A(512)(x)           #101, Connect to Detector 3
  265.     detect1 = RepConv(256)(route5)
  266.     detect2 = RepConv(512)(route6)
  267.     detect3 = RepConv(1024)(route7)
  268.     output1 = IDetect(256, 85, 3, 80)(detect1)
  269.     output2 = IDetect(512, 85, 3, 40)(detect2)
  270.     output3 = IDetect(1024, 85, 3, 20)(detect3)
  271.     output = l.Concatenate(axis=-2)([output1, output2, output3])
  272.     output = l.Activation('linear', dtype='float32')(output)
  273.     model = keras.Model(inputs=inputs, outputs=output, name="yolov7_model")
  274.     return model

损失函数的定义

YOLOv7对损失的定义可以见我另一篇文章的解读解读YOLO v7的代码(三)损失函数_gzroy的博客-CSDN博客

具体的定义在loss.py文件，我也是按照Yolov7的代码处理方式来进行tensorflow的改写，并且用了tf_function的封装来提高计算的效率， 代码如下：

  1. import tensorflow as tf
  2. import math
  3. from test1 import batch_size, na, nl, img_size, stride, balance
  4. from test1 import loss_box, loss_obj, loss_cls
  5. from test1 import batch_no_constant, anchor_no_constant, anchors_reshape, anchor_t, anchors_constant, layer_no_constant
  6. from test1 import val_batch_no_constant, val_layer_no_constant
  7. from util import *
  8. from params import *
  9.  
  10. #In param: 
  11. #    p - predictions of the model, list of three detection level.
  12. #    labels - the label of the object, dimension [batch, boxnum, 5(class, xywh)]
  13. #Out param:
  14. #    results - list of the suggest positive samples for three detection level. 
  15. #        dimension for each element: [sample_number, 5(batch_no, anch_no, x, y, class)]
  16. #    anch - list of the anchor wh ratio for the positive samples
  17. #        dimension for each element: [sample_number, anchor_w, anchor_h]
  18. @tf.function(
  19.     input_signature=(
  20.         [tf.TensorSpec(shape=[batch_size, None, 5], dtype=tf.float32)]
  21.     )
  22. )
  23. def tf_find_3_positive(labels):
  24.     batch_no = tf.zeros_like(labels)[...,0:1] + batch_no_constant
  25.     targets = tf.concat((batch_no, labels), axis=-1)    #targets dim [batch,box_num,6]
  26.     targets = tf.reshape(targets, [batch_size, 1, -1, 6])   #targets dim [batch,1,box_num,6]
  27.     targets = tf.tile(targets, [1,na,1,1])
  28.     anchor_no = anchor_no_constant + tf.reshape(tf.zeros_like(batch_no), [batch_size, 1, -1, 1])
  29.     targets = tf.concat([targets,anchor_no], axis=-1)   #targets dim [batch,na,box_num,7(batch_no, cls, xywh, anchor_no)]
  30.  
  31.     g = 0.5  # bias
  32.     offsets = tf.expand_dims(tf.constant([[0.,0.], [-1.,0.], [0.,-1.], [1.,0.], [0.,1.]]), axis=0)  #offset dim [1,5,2]
  33.  
  34.     gain = tf.constant([[1.,1.,80.,80.,80.,80.,1.], [1.,1.,40.,40.,40.,40.,1.], [1.,1.,20.,20.,20.,20.,1.]])
  35.  
  36.     results = tf.TensorArray(tf.int32, size=nl, dynamic_size=False)
  37.     anch = tf.TensorArray(tf.float32, size=nl, dynamic_size=False)
  38.  
  39.     for i in tf.range(nl):
  40.         t = targets * tf.gather(gain, i)
  41.         r = t[..., 4:6] / tf.gather(anchors_reshape, i)
  42.         r_reciprocal = tf.math.reciprocal_no_nan(r)      #1/r
  43.         r_max = tf.reduce_max(tf.math.maximum(r, r_reciprocal), axis=-1)
  44.         mask_t = tf.logical_and(r_max<anchor_t, r_max>0)
  45.         t = t[mask_t]
  46.         # Offsets
  47.         gxy = t[:, 2:4]  # grid xy
  48.         #gxi = gain[[2, 3]] - gxy  # inverse    
  49.         gxi = tf.gather(gain, i)[2:4] - gxy
  50.         mask_xy = tf.concat([
  51.             tf.ones([tf.shape(t)[0], 1], dtype=tf.bool),
  52.             ((gxy % 1. < g) & (gxy > 1.)),
  53.             ((gxi % 1. < g) & (gxi > 1.))
  54.         ], axis=1)
  55.         t = tf.repeat(tf.expand_dims(t, axis=1), 5, axis=1)[mask_xy]
  56.         offsets_xy = (tf.expand_dims(tf.zeros_like(gxy, dtype=tf.float32), axis=1) + offsets)[mask_xy]
  57.         xy = t[...,2:4] + offsets_xy
  58.         from_which_layer = tf.ones_like(t[...,0:1]) * tf.dtypes.cast(i, tf.float32)
  59.         results = results.write(i, tf.dtypes.cast(tf.concat([t[...,0:1], t[...,-1:], xy[...,1:2], xy[...,0:1], t[...,1:2], from_which_layer], axis=-1), tf.int32))
  60.         anch = anch.write(i, tf.gather(tf.gather(anchors_constant, i), tf.dtypes.cast(t[...,-1], tf.int32)))
  61.     return results.concat(), anch.concat()
  62.  
  63. @tf.function(
  64.     input_signature=([
  65.         tf.TensorSpec(shape=[None, 4], dtype=tf.float32),
  66.         tf.TensorSpec(shape=[None, 4], dtype=tf.float32)
  67.     ])
  68. )
  69. def box_iou(box1, box2):
  70.     area1 = (box1[:,2]-box1[:,0])*(box1[:,3]-box1[:,1])
  71.     area2 = (box2[:,2]-box2[:,0])*(box2[:,3]-box2[:,1])
  72.  
  73.     intersect_wh = tf.math.minimum(box1[:,None,2:], box2[:,2:]) - tf.math.maximum(box1[:,None,:2], box2[:,:2])
  74.     intersect_wh = tf.clip_by_value(intersect_wh, clip_value_min=0, clip_value_max=img_size)
  75.     intersect_area = intersect_wh[...,0]*intersect_wh[...,1]
  76.     
  77.     iou = intersect_area/(area1[:,None]+area2-intersect_area)
  78.     return iou
  79.  
  80. @tf.function(
  81.     input_signature=([
  82.         tf.TensorSpec(shape=[None, 4], dtype=tf.float32),
  83.         tf.TensorSpec(shape=[None, 4], dtype=tf.float32)
  84.     ])
  85. )
  86. def bbox_ciou(box1, box2):
  87.     eps=1e-7
  88.     b1_x1, b1_x2 = box1[:,0]-box1[:,2]/2, box1[:,0]+box1[:,2]/2
  89.     b1_y1, b1_y2 = box1[:,1]-box1[:,3]/2, box1[:,1]+box1[:,3]/2
  90.     b2_x1, b2_x2 = box2[:,0]-box2[:,2]/2, box2[:,0]+box2[:,2]/2
  91.     b2_y1, b2_y2 = box2[:,1]-box2[:,3]/2, box2[:,1]+box2[:,3]/2
  92.     
  93.     # Intersection area
  94.     inter = tf.clip_by_value(
  95.         tf.math.minimum(b1_x2, b2_x2) - tf.math.maximum(b1_x1, b2_x1), 
  96.         clip_value_min=0, 
  97.         clip_value_max=tf.float32.max) * tf.clip_by_value(
  98.         tf.math.minimum(b1_y2, b2_y2) - tf.math.maximum(b1_y1, b2_y1), 
  99.         clip_value_min=0, 
  100.         clip_value_max=tf.float32.max)
  101.     
  102.     # Union Area
  103.     w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
  104.     w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
  105.     union = w1 * h1 + w2 * h2 - inter + eps
  106.  
  107.     iou = inter / union
  108.     
  109.     cw = tf.math.maximum(b1_x2, b2_x2) - tf.math.minimum(b1_x1, b2_x1)  # convex (smallest enclosing box) width
  110.     ch = tf.math.maximum(b1_y2, b2_y2) - tf.math.minimum(b1_y1, b2_y1)  # convex height
  111.     
  112.     c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
  113.     rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
  114.         (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
  115.     
  116.     v = (4 / math.pi ** 2) * tf.math.pow(tf.math.atan(w2 / (h2 + eps)) - tf.math.atan(w1 / (h1 + eps)), 2)
  117.     alpha = v / (v - iou + (1 + eps))
  118.     return iou - (rho2 / c2 + v * alpha)
  119.  
  120. @tf.function(
  121.     input_signature=([
  122.         tf.TensorSpec(shape=[batch_size, na, None, 85], dtype=tf.float32),
  123.         tf.TensorSpec(shape=[batch_size, None, 5], dtype=tf.float32)
  124.     ])
  125. )
  126. def tf_build_targets(p, labels):
  127.     results, anch = tf_find_3_positive(labels)
  128.  
  129.     #stride = tf.constant([8., 16., 32.])
  130.     grids = tf.dtypes.cast(img_size/stride, tf.int32)
  131.  
  132.     pxyxys = tf.TensorArray(tf.float32, size=nl, dynamic_size=False)
  133.     p_obj = tf.TensorArray(tf.float32, size=nl, dynamic_size=True, element_shape=[None, 1])
  134.     p_cls = tf.TensorArray(tf.float32, size=nl, dynamic_size=False)
  135.     all_idx = tf.TensorArray(tf.int32, size=nl, dynamic_size=False)
  136.     from_which_layer = tf.TensorArray(tf.int32, size=nl, dynamic_size=False)
  137.     all_anch = tf.TensorArray(tf.float32, size=nl, dynamic_size=False)
  138.     
  139.     matching_idxs = tf.TensorArray(tf.int32, size=batch_size, dynamic_size=False)
  140.     matching_targets = tf.TensorArray(tf.float32, size=batch_size, dynamic_size=False)
  141.     matching_anchs = tf.TensorArray(tf.float32, size=batch_size, dynamic_size=False)
  142.     matching_layers = tf.TensorArray(tf.int32, size=batch_size, dynamic_size=False)
  143.  
  144.     for i in tf.range(nl):
  145.         idx_mask = results[...,-1]==i
  146.         idx = tf.boolean_mask(results, idx_mask)
  147.         layer_mask = layer_no_constant[...,0]==i
  148.         grid_no = tf.gather(grids, i)
  149.         pl = tf.boolean_mask(p, layer_mask)
  150.         pl = tf.reshape(pl, [batch_size, na, grid_no, grid_no, -1])
  151.         pi = tf.gather_nd(pl, idx[...,0:4])
  152.         anchors_p = tf.boolean_mask(anch, idx_mask)
  153.         p_obj = p_obj.write(i, pi[...,4:5])
  154.         p_cls = p_cls.write(i, pi[...,5:])
  155.         gij = tf.dtypes.cast(tf.concat([idx[...,3:4], idx[...,2:3]], axis=-1), tf.float32)
  156.         pxy = (tf.math.sigmoid(pi[...,:2])*2-0.5+gij)*tf.dtypes.cast(tf.gather(stride, i), tf.float32)
  157.         pwh = (tf.math.sigmoid(pi[...,2:4])*2)**2*anchors_p*tf.dtypes.cast(tf.gather(stride, i), tf.float32)
  158.         pxywh = tf.concat([pxy, pwh], axis=-1)
  159.         pxyxy = xywh2xyxy(pxywh)
  160.         pxyxys = pxyxys.write(i, pxyxy)
  161.         all_idx = all_idx.write(i, idx[...,0:4])
  162.         from_which_layer = from_which_layer.write(i, idx[..., -1:])
  163.         all_anch = all_anch.write(i, tf.boolean_mask(anch, idx_mask))
  164.  
  165.     pxyxys = pxyxys.concat()
  166.     p_obj = p_obj.concat()
  167.     p_cls = p_cls.concat()
  168.     all_idx = all_idx.concat()
  169.     from_which_layer = from_which_layer.concat()
  170.     all_anch = all_anch.concat()
  171.  
  172.     for i in tf.range(batch_size):
  173.         batch_mask = all_idx[...,0]==i
  174.         if tf.math.reduce_sum(tf.dtypes.cast(batch_mask, tf.int32)) > 0:
  175.             pxyxy_i = tf.boolean_mask(pxyxys, batch_mask)
  176.             target_mask = labels[i][...,3]>0
  177.             target = tf.boolean_mask(labels[i], target_mask)
  178.             txywh = target[...,1:] * img_size
  179.             txyxy = xywh2xyxy(txywh)
  180.             pair_wise_iou = box_iou(txyxy, pxyxy_i)
  181.             pair_wise_iou_loss = -tf.math.log(pair_wise_iou + 1e-8)
  182.  
  183.             top_k, _ = tf.math.top_k(pair_wise_iou, tf.math.minimum(10, tf.shape(pair_wise_iou)[1]))
  184.             dynamic_ks = tf.clip_by_value(
  185.                 tf.dtypes.cast(tf.math.reduce_sum(top_k, axis=-1), tf.int32),
  186.                 clip_value_min=1, 
  187.                 clip_value_max=10)
  188.  
  189.             gt_cls_per_image = tf.tile(
  190.                 tf.expand_dims(
  191.                     tf.one_hot(
  192.                         tf.dtypes.cast(target[...,0], tf.int32), nc),
  193.                     axis = 1),
  194.                 [1,tf.shape(pxyxy_i)[0],1])
  195.  
  196.             num_gt = tf.shape(target)[0]
  197.             cls_preds_ = (
  198.                 tf.math.sigmoid(tf.tile(tf.expand_dims(tf.boolean_mask(p_cls, batch_mask), 0), [num_gt, 1, 1])) *
  199.                 tf.math.sigmoid(tf.tile(tf.expand_dims(tf.boolean_mask(p_obj, batch_mask), 0), [num_gt, 1, 1])))    #dimension [labels_number, positive_targets_number, 80]
  200.             y = tf.math.sqrt(cls_preds_)
  201.             pair_wise_cls_loss = tf.math.reduce_sum(
  202.                 tf.nn.sigmoid_cross_entropy_with_logits(
  203.                     labels = gt_cls_per_image,
  204.                     logits = tf.math.log(y/(1-y))),
  205.                 axis = -1)
  206.  
  207.             cost = (
  208.                 pair_wise_cls_loss
  209.                 + 3.0 * pair_wise_iou_loss
  210.             )
  211.  
  212.             matching_matrix = tf.zeros_like(cost)      #dimension [labels_number, positive_targets_number]
  213.  
  214.             matching_idx = tf.TensorArray(tf.int64, size=0, dynamic_size=True)
  215.             for gt_idx in tf.range(num_gt):
  216.                 _, pos_idx = tf.math.top_k(
  217.                     -cost[gt_idx], k=dynamic_ks[gt_idx], sorted=True)
  218.                 X,Y = tf.meshgrid(gt_idx, pos_idx)
  219.                 matching_idx = matching_idx.write(gt_idx, tf.dtypes.cast(tf.concat([X,Y], axis=-1), tf.int64))
  220.  
  221.             matching_idx = matching_idx.concat()
  222.             '''
  223.             matching_matrix = tf.scatter_nd(
  224.                 matching_idx, 
  225.                 tf.ones(tf.shape(matching_idx)[0]), 
  226.                 tf.dtypes.cast(tf.shape(cost), tf.int64))
  227.             '''
  228.             matching_matrix = tf.sparse.to_dense(
  229.                 tf.sparse.reorder(
  230.                     tf.sparse.SparseTensor(
  231.                         indices=tf.dtypes.cast(matching_idx, tf.int64), 
  232.                         values=tf.ones(tf.shape(matching_idx)[0]), 
  233.                         dense_shape=tf.dtypes.cast(tf.shape(cost), tf.int64))
  234.                 )
  235.             )
  236.  
  237.             anchor_matching_gt = tf.reduce_sum(matching_matrix, axis=0)    #dimension [positive_targets_number]
  238.             mask_1 = anchor_matching_gt>1     #it means one target match to several ground truths
  239.  
  240.             if tf.reduce_sum(tf.dtypes.cast(mask_1, tf.int32)) > 0:   #There is at least one positive target that predict several ground truth  
  241.                 #Get the lowest cost of the serveral ground truth of the target
  242.                 #For example, there are 100 targets and 10 ground truths.
  243.                 #The #5 target match to the #2 and #3 ground truth, the related cost are 10 for #2 and 20 for #3
  244.                 #Then it will select #2 gound truth for the #5 target.
  245.                 #mask_1 dimension [positive_targets_number]
  246.                 #tf.boolean_mask(cost, mask_1, axis=1), dimension [ground_truth_numer, targets_predict_sevearl_GT_number]
  247.                 cost_argmin = tf.math.argmin(
  248.                     tf.boolean_mask(cost, mask_1, axis=1), axis=0)  #in above example, the cost_argmin is [2]
  249.                 m = tf.dtypes.cast(mask_1, tf.float32)
  250.                 _, target_indices = tf.math.top_k(
  251.                     m, 
  252.                     k=tf.dtypes.cast(tf.math.reduce_sum(m), tf.int32))  #in above example, the target_indices is [5]
  253.                 #So will set the index [2,5] of matching_matrix to 1, and set the other elements of [:,5] to 0
  254.                 target_matching_gt_indices = tf.concat(
  255.                     [tf.reshape(tf.dtypes.cast(cost_argmin, tf.int32), [-1,1]), tf.reshape(target_indices, [-1,1])], 
  256.                     axis=1)          
  257.                 matching_matrix = tf.multiply(
  258.                     matching_matrix,
  259.                     tf.repeat(tf.reshape(tf.dtypes.cast(anchor_matching_gt<=1, tf.float32), [1,-1]), tf.shape(cost)[0], axis=0))
  260.                 target_value = tf.sparse.to_dense(
  261.                     tf.sparse.reorder(
  262.                         tf.sparse.SparseTensor(
  263.                             indices=tf.dtypes.cast(target_matching_gt_indices, tf.int64),
  264.                             values=tf.ones(tf.shape(target_matching_gt_indices)[0]),
  265.                             dense_shape=tf.dtypes.cast(tf.shape(matching_matrix), tf.int64)
  266.                         )
  267.                     )
  268.                 )
  269.                 matching_matrix = tf.add(matching_matrix, target_value)
  270.  
  271.             fg_mask_inboxes = tf.math.reduce_sum(matching_matrix, axis=0)>0.  #The mask for the targets that will use to predict
  272.             if tf.shape(tf.boolean_mask(matching_matrix, fg_mask_inboxes, axis=1))[0]>0:
  273.                 matched_gt_inds = tf.math.argmax(tf.boolean_mask(matching_matrix, fg_mask_inboxes, axis=1), axis=0)  #Get the related gt number for the target
  274.  
  275.                 all_idx_i = tf.boolean_mask(tf.boolean_mask(all_idx, batch_mask), fg_mask_inboxes)
  276.                 from_which_layer_i = tf.boolean_mask(tf.boolean_mask(from_which_layer, batch_mask), fg_mask_inboxes)
  277.                 all_anch_i = tf.boolean_mask(tf.boolean_mask(all_anch, batch_mask), fg_mask_inboxes)
  278.  
  279.                 matching_idxs = matching_idxs.write(i, all_idx_i)
  280.                 matching_layers = matching_layers.write(i, from_which_layer_i)
  281.                 matching_anchs = matching_anchs.write(i, all_anch_i )
  282.                 matching_targets = matching_targets.write(i, tf.gather(target, matched_gt_inds))
  283.             else:
  284.                 matching_idxs = matching_idxs.write(i, tf.constant([[-1,-1,-1,-1]], dtype=tf.int32))
  285.                 matching_layers = matching_layers.write(i, tf.constant([[-1]], dtype=tf.int32))
  286.                 matching_anchs = matching_anchs.write(i, tf.constant([[-1, -1]], dtype=tf.float32))
  287.                 matching_targets = matching_targets.write(i, tf.constant([[-1, -1, -1, -1, -1]], dtype=tf.float32))                                    
  288.         
  289.         else:
  290.             matching_idxs = matching_idxs.write(i, tf.constant([[-1,-1,-1,-1]], dtype=tf.int32))
  291.             matching_layers = matching_layers.write(i, tf.constant([[-1]], dtype=tf.int32))
  292.             matching_anchs = matching_anchs.write(i, tf.constant([[-1, -1]], dtype=tf.float32))
  293.             matching_targets = matching_targets.write(i, tf.constant([[-1, -1, -1, -1, -1]], dtype=tf.float32))
  294.         
  295.     matching_idxs = matching_idxs.concat()
  296.     matching_layers = matching_layers.concat()
  297.     matching_anchs = matching_anchs.concat()
  298.     matching_targets = matching_targets.concat()
  299.     filter_mask = matching_idxs[:,0]!=-1
  300.     matching_idxs = tf.boolean_mask(matching_idxs, filter_mask)
  301.     matching_layers = tf.boolean_mask(matching_layers, filter_mask)
  302.     matching_anchs = tf.boolean_mask(matching_anchs, filter_mask)
  303.     matching_targets = tf.boolean_mask(matching_targets, filter_mask)
  304.     
  305.     #return pxyxys, all_idx, matching_idx, matching_matrix, all_idx_i, cost, pair_wise_iou, from_which_layer_i
  306.     return matching_idxs, matching_layers, matching_anchs, matching_targets
  307.  
  308. @tf.function(
  309.     input_signature=([
  310.         tf.TensorSpec(shape=[batch_size, na, None, 85], dtype=tf.float32),
  311.         tf.TensorSpec(shape=[batch_size, None, 5], dtype=tf.float32)
  312.     ])
  313. )
  314. def tf_loss_func(p, labels):
  315.     matching_idxs, matching_layers, matching_anchs, matching_targets = tf_build_targets(p, labels)
  316.     lcls, lbox, lobj = tf.zeros(1), tf.zeros(1), tf.zeros(1)
  317.     
  318.     grids = img_size//stride
  319.     for i in tf.range(nl):
  320.         layer_mask = layer_no_constant[...,0]==i
  321.         grid = tf.gather(grids, i)
  322.         pi = tf.reshape(tf.boolean_mask(p, layer_mask), [batch_size, na, grid, grid, -1])
  323.         matching_layer_mask = matching_layers[:,0]==i
  324.         if tf.reduce_sum(tf.dtypes.cast(matching_layer_mask, tf.int32))==0:
  325.             continue
  326.         m_idxs = tf.boolean_mask(matching_idxs, matching_layer_mask)
  327.         if tf.shape(m_idxs)[0]==0:
  328.             continue
  329.         m_targets = tf.boolean_mask(matching_targets, matching_layer_mask)
  330.         m_anchs = tf.boolean_mask(matching_anchs, matching_layer_mask)
  331.         ps = tf.gather_nd(pi, m_idxs)
  332.         pxy = tf.math.sigmoid(ps[:,:2])*2-0.5
  333.         pwh = (tf.math.sigmoid(ps[:,2:4])*2)**2*m_anchs
  334.         pbox = tf.concat([pxy,pwh], axis=-1)
  335.         #selected_tbox = tf.gather_nd(labels, matching_targets[i])[:, 1:]
  336.         selected_tbox = m_targets[:, 1:]
  337.         selected_tbox = tf.multiply(selected_tbox, tf.dtypes.cast(grid, tf.float32))
  338.         tbox_grid = tf.concat([
  339.             tf.dtypes.cast(m_idxs[:,3:4], tf.float32),
  340.             tf.dtypes.cast(m_idxs[:,2:3], tf.float32),
  341.             tf.zeros((tf.shape(m_idxs)[0],2))], 
  342.             axis=-1)
  343.         selected_tbox = tf.subtract(selected_tbox, tbox_grid)
  344.         iou = bbox_ciou(pbox, selected_tbox)
  345.         lbox += tf.math.reduce_mean(1.0 - iou)  # iou loss
  346.  
  347.         # Objectness
  348.         tobj = tf.sparse.to_dense(
  349.             tf.sparse.reorder(
  350.                 tf.sparse.SparseTensor(
  351.                     indices = tf.dtypes.cast(m_idxs, tf.int64),
  352.                     values = (1.0 - gr) + gr * tf.clip_by_value(tf.stop_gradient(iou), clip_value_min=0, clip_value_max=tf.float32.max),
  353.                     dense_shape = tf.dtypes.cast(tf.shape(pi[..., 0]), tf.int64)
  354.                 )
  355.             ), validate_indices=False
  356.         )
  357.  
  358.         # Classification
  359.  
  360.         tcls = tf.one_hot(
  361.             indices = tf.dtypes.cast(m_targets[:,0], tf.int32),
  362.             depth = 80,
  363.             dtype = tf.float32
  364.         )
  365.         
  366.         lcls += tf.math.reduce_mean(
  367.             tf.nn.sigmoid_cross_entropy_with_logits(
  368.                 labels = tcls,
  369.                 logits = ps[:, 5:]
  370.             )
  371.         )
  372.         '''
  373.         lcls += tf.math.reduce_mean(
  374.             tf.nn.sparse_softmax_cross_entropy_with_logits(
  375.                 labels = tf.dtypes.cast(m_targets[:,0], tf.int32),
  376.                 logits = ps[:, 5:]
  377.             )    
  378.         )
  379.         '''
  380.         obji = tf.math.reduce_mean(
  381.             tf.nn.sigmoid_cross_entropy_with_logits(
  382.                 labels = tobj,
  383.                 logits = pi[..., 4]
  384.             )
  385.         )
  386.  
  387.         lobj += obji * tf.gather(balance, i) 
  388.         
  389.     lbox *= loss_box
  390.     lobj *= loss_obj
  391.     lcls *= loss_cls
  392.  
  393.     loss = (lbox + lobj + lcls) * batch_size
  394.  
  395.     return loss
  396.  
  397. @tf.function(
  398.     input_signature=([
  399.         tf.TensorSpec(shape=[None, na, 8400, 85], dtype=tf.float32),
  400.         tf.TensorSpec(shape=[None, None, 5], dtype=tf.float32),
  401.         tf.TensorSpec(shape=[None, 2], dtype=tf.int32),
  402.         tf.TensorSpec(shape=[None], dtype=tf.int32),
  403.     ])
  404. )
  405. def tf_predict_func(predictions, labels, imgs_hw, imgs_id):
  406.     grids = img_size // stride
  407.     batch_size = tf.shape(predictions)[0]
  408.     confidence_threshold = 0.2
  409.     probabilty_threshold = 0.8
  410.     all_predict_result = tf.TensorArray(tf.float32, size=nl, dynamic_size=False) 
  411.     boxes_result = tf.TensorArray(tf.float32, size=0, dynamic_size=True) 
  412.     imgs_info = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
  413.     for i in tf.range(nl):
  414.         grid = tf.gather(grids, i)
  415.         grid_x, grid_y = tf.meshgrid(tf.range(grid, dtype=tf.float32), tf.range(grid, dtype=tf.float32))
  416.         grid_x = tf.reshape(grid_x, [-1, 1])
  417.         grid_y = tf.reshape(grid_y, [-1, 1])
  418.         #grid_xy = tf.concat([grid_y, grid_x], axis=-1)
  419.         grid_xy = tf.concat([grid_x, grid_y], axis=-1)
  420.         grid_xy = tf.reshape(grid_xy, [1,1,-1,2])
  421.         layer_mask = val_layer_no_constant[...,0]==i
  422.         #grid = tf.gather(grids, i)
  423.         predict_layer = tf.boolean_mask(predictions, layer_mask)
  424.         predict_layer = tf.reshape(predict_layer, [batch_size, na, -1, 85])
  425.         predict_conf = tf.math.sigmoid(predict_layer[...,4:5])
  426.         predict_xy = (tf.math.sigmoid(predict_layer[...,:2])*2-0.5 + \
  427.             tf.dtypes.cast(grid_xy,tf.float32))*tf.dtypes.cast(tf.gather(stride, i), tf.float32)
  428.         predict_wh = (tf.math.sigmoid(predict_layer[...,2:4])*2)**2*\
  429.             tf.reshape(tf.gather(anchors_constant,i), [1,na,1,2])*\
  430.             tf.dtypes.cast(tf.gather(stride, i), tf.float32)
  431.         predict_xywh = tf.concat([predict_xy, predict_wh], axis=-1)
  432.         predict_xyxy = xywh2xyxy(predict_xywh)
  433.         predict_cls = tf.reshape(tf.argmax(predict_layer[...,5:], axis=-1), [batch_size, na, -1, 1])
  434.         predict_cls = tf.dtypes.cast(predict_cls, tf.float32)
  435.         predict_proba = tf.nn.sigmoid(
  436.             tf.reduce_max(
  437.                 predict_layer[...,5:], axis=-1, keepdims=True
  438.             )
  439.         )
  440.         batch_no = tf.expand_dims(tf.tile(tf.gather(val_batch_no_constant, tf.range(batch_size)), [1,na,grid*grid]), -1)
  441.         predict_result = tf.concat([batch_no, predict_conf, predict_xyxy, predict_cls, predict_proba], axis=-1)
  442.         mask = tf.math.logical_and(
  443.             predict_result[...,1]>=confidence_threshold,
  444.             predict_result[...,-1]>=probabilty_threshold
  445.         )
  446.         predict_result = tf.boolean_mask(predict_result, mask)
  447.         #tf.print(tf.shape(predict_result))
  448.         if tf.shape(predict_result)[0] > 0:
  449.             all_predict_result = all_predict_result.write(i, predict_result)
  450.             #tf.print(tf.shape(predict_result))
  451.         else:
  452.             all_predict_result = all_predict_result.write(i, tf.zeros(shape=[1,8]))
  453.     all_predict_result = all_predict_result.concat()
  454.     #return all_predict_result
  455.           
  456.     for i in tf.range(batch_size):
  457.         batch_mask = tf.math.logical_and(
  458.             all_predict_result[...,0]==tf.dtypes.cast(i, tf.float32),
  459.             all_predict_result[...,1]>0
  460.         )
  461.         predict_true_box = tf.boolean_mask(all_predict_result, batch_mask)
  462.         if tf.shape(predict_true_box)[0]==0:
  463.             continue
  464.         original_hw = tf.dtypes.cast(tf.gather(imgs_hw, i), tf.float32)
  465.         ratio = tf.dtypes.cast(tf.reduce_max(original_hw/img_size), tf.float32)
  466.         predict_classes, _ = tf.unique(predict_true_box[:,6])
  467.         #predict_classes_list = tf.unstack(predict_classes)
  468.         #for class_id in predict_classes_list:
  469.         for j in tf.range(tf.shape(predict_classes)[0]):
  470.             #class_mask = tf.math.equal(predict_true_box[:, 6], class_id)
  471.             class_mask = tf.math.equal(predict_true_box[:, 6], tf.gather(predict_classes, j))
  472.             predict_true_box_class = tf.boolean_mask(predict_true_box, class_mask)
  473.             predict_true_box_xy = predict_true_box_class[:, 2:6]
  474.             predict_true_box_score = predict_true_box_class[:, 7]*predict_true_box_class[:, 1]
  475.             #predict_true_box_score = predict_true_box_class[:, 1]
  476.             selected_indices = tf.image.non_max_suppression(
  477.                 predict_true_box_xy,
  478.                 predict_true_box_score,
  479.                 100,
  480.                 iou_threshold=0.2
  481.                 #score_threshold=confidence_threshold
  482.             )
  483.             #Shape [box_num, 7]
  484.             selected_boxes = tf.gather(predict_true_box_class, selected_indices) 
  485.             #boxes_result = boxes_result.write(boxes_result.size(), selected_boxes)
  486.             boxes_xyxy = selected_boxes[:,2:6]*ratio
  487.             boxes_x1 = tf.clip_by_value(boxes_xyxy[:,0:1], 0., original_hw[1])
  488.             boxes_x2 = tf.clip_by_value(boxes_xyxy[:,2:3], 0., original_hw[1])
  489.             boxes_y1 = tf.clip_by_value(boxes_xyxy[:,1:2], 0., original_hw[0])
  490.             boxes_y2 = tf.clip_by_value(boxes_xyxy[:,3:4], 0., original_hw[0])
  491.             boxes_w = boxes_x2 - boxes_x1
  492.             boxes_h = boxes_y2 - boxes_y1
  493.             boxes = tf.concat([selected_boxes[:,0:2], boxes_x1, boxes_y1, boxes_w, boxes_h, selected_boxes[:,6:8]], axis=-1)
  494.             boxes_result = boxes_result.write(boxes_result.size(), boxes)
  495.         img_id = tf.gather(imgs_id, i)
  496.         imgs_info = imgs_info.write(imgs_info.size(), tf.reshape(tf.stack([i, img_id]), [-1,2]))
  497.     if boxes_result.size()==0:
  498.         boxes_result = boxes_result.write(0, tf.zeros(shape=[1,8]))
  499.     if imgs_info.size()==0:
  500.         imgs_info = imgs_info.write(0, tf.dtypes.cast(tf.zeros(shape=[1,2]), tf.int32))
  501.  
  502.     return boxes_result.concat(), imgs_info.concat()

训练与验证

最后就是对模型进行训练和验证了，这里也是按照YOLOv7的实现方式来进行训练，验证的时候是采用pycocotools工具来进行mAP的计算。具体可以参见train.py文件

因为模型是对640*640大小的图像进行训练，对GPU的显存要求很大。在我本地的2080Ti显卡，11G内存的情况下，开启混合精度，只能设置Batch size为8，训练效果不是很理想。为此我在autodl平台租用了一个V100的32G显存的GPU来进行测试（价格是每小时2.28元），Batch size设置为32。感觉Batch size对模型的训练效果还是有比较大的影响的。最终经过了20多个epoch的训练，每个Epoch大概要训练1个小时多一点，大概花费了1天的时间，结果如下：

  1.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.270
  2.  Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.411
  3.  Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.289
  4.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.162
  5.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.302
  6.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.334
  7.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.268
  8.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.476
  9.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.528
  10.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.338
  11.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.576
  12.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.661

以下是对验证集的一些图片的预测结果，

按照Yolov7论文的描述，训练了300个epoch之后，mAP all能达到60%，继续训练可以进一步提高准确率，不过限于时间和资源，我就暂时训练到这个地步。

最后，我的源码都放在了Github的仓库，GitHub - gzroy/yolov7_tf2: Yolov7 implementation on tensorflow 2.x

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家自动化/article/detail/295652
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号