训练测试验证tflite.

detect.py

@@ -16,13 +16,12 @@
 
 def main():
     # model_path = "h5模型路径, 默认在 ./logs/yolov5-tf-300.h5"
-    # model_path = "./logs_bn_momentum0.75/yolov5-tf-300.h5"
     model_path = "./logs/yolov5s-best.h5"
     # image_path = "提供你要测试的图片路径"
-    # image_path = "./data/tmp/traffic_road.jpg"
+    image_path = "./data/tmp/Cats_Test49.jpg"
     # image_path = "./data/coco_2017_val_images/289343.jpg"
     # image_path = "./data/cat_dog_face_data/JPEGImages/Cats_Test849.jpg"
-    image_path = "./data/cat_dog_face_data/JPEGImages/Cats_Test360.jpg"
+    # image_path = "./data/cat_dog_face_data/JPEGImages/Cats_Test214.jpg"
     image = cv2.imread(image_path)
     # 可以选择 ['5l', '5s', '5m', '5x'], 跟随训练
     yolov5_type = "5s"

train.py

@@ -12,6 +12,7 @@
 from yolo import Yolo
 from loss import ComputeLoss
 from val import val
+from layers import nms
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -159,7 +160,7 @@ def main():
                 # pred, 同样只拿第一个batch的pred
                 pred_img = gt_imgs[random_one].copy() * 255
                 yolo_head_output = yolo.yolo_head(yolo_preds, is_training=False)
-                nms_output = yolo.nms(yolo_head_output.numpy(), iou_thres=0.3)
+                nms_output = nms(image_shape, yolo_head_output.numpy(), iou_thres=0.3)
                 if len(nms_output) == batch_size:
                     nms_output = nms_output[random_one]
                     for box_obj_cls in nms_output:

val.py

@@ -7,6 +7,7 @@
 from yolo import Yolo
 import tensorflow as tf
 from data.generate_coco_data import CoCoDataGenrator
+from layers import nms
 from data.visual_ops import draw_bounding_box
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "1"
@@ -110,7 +111,7 @@ def val(model, val_data_generator, classes, desc='val'):
             # predictions = model.yolov5(gt_imgs / 255., training=True)
             predictions = model.yolov5.predict(gt_imgs / 255.)
             predictions = model.yolo_head(predictions, is_training=False)
-            predictions = model.nms(predictions.numpy())
+            predictions = nms(model.image_shape, predictions.numpy())
         else:
             predictions = model.predict(gt_imgs, image_need_resize=False, resize_to_origin=False)
 

yolo.py

@@ -10,69 +10,11 @@
 from yolov5x import Yolov5x
 from yolov5m import Yolov5m
 from yolov5s import Yolov5s
-from data.generate_coco_data import CoCoDataGenrator
-from data.visual_ops import draw_bounding_box
+from layers import nms, YoloHead
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 
-class YoloHead(tf.keras.layers.Layer):
-    def __init__(self, image_shape, num_class, is_training, strides, anchors, anchors_masks):
-        super(YoloHead, self).__init__()
-        self.image_shape = image_shape
-        self.num_class = num_class
-        self.is_training = is_training
-        self.strides = strides
-        self.anchors = anchors
-        self.anchors_masks = anchors_masks
-        self.grid = []
-        self.anchor_grid = []
-        for i, stride in enumerate(strides):
-            grid, anchor_grid = self._make_grid(self.image_shape[0] // stride, self.image_shape[1] // stride, i)
-            self.grid.append(grid)
-            self.anchor_grid.append(anchor_grid)
-
-    def call(self, inputs, *args, **kwargs):
-        detect_res = []
-        for i, pred in enumerate(inputs):
-            if not self.is_training:
-                # f_shape = tf.shape(pred)
-                # if len(self.grid) < self.anchor_masks.shape[0]:
-                #     grid, anchor_grid = self._make_grid(f_shape[1], f_shape[2], i)
-                #     self.grid.append(grid)
-                #     self.anchor_grid.append(anchor_grid)
-                # 这里把输出的值域从[0,1]调整到[0, image_shape]
-                pred_xy = (tf.sigmoid(pred[..., 0:2]) * 2. - 0.5 + self.grid[i]) * self.strides[i]
-                pred_wh = (tf.sigmoid(pred[..., 2:4]) * 2) ** 2 * self.anchor_grid[i]
-                # print(self.grid)
-                pred_obj = tf.sigmoid(pred[..., 4:5])
-                pred_cls = tf.keras.layers.Softmax()(pred[..., 5:])
-                cur_layer_pred_res = tf.keras.layers.Concatenate(axis=-1)([pred_xy, pred_wh, pred_obj, pred_cls])
-
-                # cur_layer_pred_res = tf.reshape(cur_layer_pred_res, [self.batch_size, -1, self.num_class + 5])
-                cur_layer_pred_res = tf.keras.layers.Reshape([-1, self.num_class + 5])(cur_layer_pred_res)
-                detect_res.append(cur_layer_pred_res)
-            else:
-                detect_res.append(pred)
-        return detect_res if self.is_training else tf.concat(detect_res, axis=1)
-
-    def _make_grid(self, h, w, i):
-        cur_layer_anchors = self.anchors[self.anchors_masks[i]] * np.array([[self.image_shape[1], self.image_shape[0]]])
-        num_anchors_per_layer = len(cur_layer_anchors)
-        yv, xv = tf.meshgrid(tf.range(h), tf.range(w))
-        grid = tf.stack((xv, yv), axis=2)
-        # 用来计算中心点的grid cell左上角坐标
-        grid = tf.tile(tf.reshape(grid, [1, h, w, 1, 2]), [1, 1, 1, num_anchors_per_layer, 1])
-        grid = tf.cast(grid, tf.float32)
-        # anchor_grid = tf.reshape(cur_layer_anchors * self.strides[i], [1, 1, 1, num_anchors_per_layer, 2])
-        anchor_grid = tf.reshape(cur_layer_anchors, [1, 1, 1, num_anchors_per_layer, 2])
-        # 用来计算宽高的anchor w/h
-        anchor_grid = tf.tile(anchor_grid, [1, h, w, 1, 1])
-        anchor_grid = tf.cast(anchor_grid, tf.float32)
-
-        return grid, anchor_grid
-
-
 class Yolo:
 
     def __init__(self,
@@ -220,59 +162,6 @@ def _make_grid(self, h, w, i):
 
         return grid, anchor_grid
 
-    def nms(self, predicts, conf_thres=0.45, iou_thres=0.2, max_det=300, max_nms=3000):
-        """ 原yolov5简化版nms, 不用multi label, 不做merge box
-
-        :param predicts:
-        :param conf_thres:
-        :param iou_thres:
-        :param max_det:
-        :return: [batch, nms_nums, (x1, y1, x2, y2, conf, cls)]
-        """
-        output = []
-
-        # 这里遍历每个batch,也就是每张图, 输出的3层预测已经做了合并处理成[batch, -1, 5+num_class]
-        for i, predict in enumerate(predicts):
-            # predict = predict.numpy()
-            # 首先只要那些目标概率大于阈值的
-            obj_mask = predict[..., 4] > conf_thres
-            predict = predict[obj_mask]
-
-            # 没有满足的数据则跳过去下一张
-            if not predict.shape[0]:
-                continue
-
-            # 类别概率乘上了目标概率, 作为最终判别概率
-            predict[:, 5:] *= predict[:, 4:5]
-
-            x1 = np.maximum(predict[:, 0] - predict[:, 2] / 2, 0)
-            y1 = np.maximum(predict[:, 1] - predict[:, 3] / 2, 0)
-            x2 = np.minimum(predict[:, 0] + predict[:, 2] / 2, self.image_shape[1])
-            y2 = np.minimum(predict[:, 1] + predict[:, 3] / 2, self.image_shape[0])
-            box = np.concatenate([x1[:, None], y1[:, None], x2[:, None], y2[:, None]], axis=-1)
-            # Detections matrix [n, (x1, y1, x2, y2, conf, cls)]
-            max_cls_ids = np.array(predict[:, 5:].argmax(axis=1), dtype=np.float32)
-            max_cls_score = predict[:, 5:].max(axis=1)
-            predict = np.concatenate([box, max_cls_score[:, None], max_cls_ids[:, None]], axis=1)[
-                np.reshape(max_cls_score > 0.1, (-1,))]
-
-            n = predict.shape[0]
-            if not n:
-                continue
-            elif n > max_nms:
-                predict = predict[predict[:, 4].argsort()[::-1][:max_nms]]
-
-            # 为每个类别乘上一个大数,box再加上这个偏移, 做nms时就可以在类内做
-            cls = predict[:, 5:6] * 4096
-            # 边框加偏移
-            boxes, scores = predict[:, :4] + cls, predict[:, 4]
-            nms_ids = tf.image.non_max_suppression(
-                boxes=boxes, scores=scores, max_output_size=max_det, iou_threshold=iou_thres)
-
-            output.append(predict[nms_ids.numpy()])
-
-        return output
-
     def build_graph(self):
         # inputs = tf.keras.layers.Input(shape=self.image_shape, batch_size=self.batch_size)
         inputs = tf.keras.layers.Input(shape=self.image_shape)
@@ -290,6 +179,7 @@ def build_graph(self):
         )(yolo_body_outputs)
         # if not self.is_training:
         #     outputs = self.nms(outputs, iou_thres=self.yolo_iou_threshold, conf_thres=self.yolo_conf_threshold)
+        # model = tf.keras.models.Model(inputs=inputs, outputs=yolo_body_outputs)
         model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
         return model
 
@@ -335,7 +225,7 @@ def predict(self, images, image_need_resize=True, resize_to_origin=True):
             # 非极大抑制, [nms_nums, (x1, y1, x2, y2, conf, cls)]
             # nms_outputs = self.nms(outputs.numpy(), iou_thres=0.3)[0]
             # print(np.max(outputs[:,:,4]),np.min(outputs[:,:,4]))
-            nms_outputs = self.nms(outputs)
+            nms_outputs = nms(self.image_shape, outputs)
             # nms_outputs = self.nms(outputs.numpy())
             # print(nms_outputs.shape)
             # if not nms_outputs.shape[0]: