add nms_class_agnostic for test

models/detectors/yolov1/build.py

@@ -25,7 +25,8 @@ def build_yolov1(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         conf_thresh = args.conf_thresh,
         nms_thresh = args.nms_thresh,
         trainable = trainable,
-        deploy = deploy
+        deploy = deploy,
+        nms_class_agnostic=args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov1/yolov1.py

@@ -19,7 +19,8 @@ class YOLOv1(nn.Module):
                  conf_thresh=0.01,
                  nms_thresh=0.5,
                  trainable=False,
-                 deploy=False):
+                 deploy=False,
+                 nms_class_agnostic :bool = False):
         super(YOLOv1, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -31,6 +32,7 @@ class YOLOv1(nn.Module):
         self.nms_thresh = nms_thresh                   # NMS阈值
         self.stride = 32                               # 网络的最大步长
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         
         # ------------------- Network Structure -------------------
         ## 主干网络
@@ -110,7 +112,7 @@ class YOLOv1(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels
 

models/detectors/yolov2/build.py

@@ -25,7 +25,8 @@ def build_yolov2(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         conf_thresh=args.conf_thresh,
         nms_thresh=args.nms_thresh,
         topk=args.topk,
-        deploy=deploy
+        deploy=deploy,
+        nms_class_agnostic=args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov2/yolov2.py

@@ -19,7 +19,8 @@ class YOLOv2(nn.Module):
                  nms_thresh=0.5,
                  topk=100,
                  trainable=False,
-                 deploy=False):
+                 deploy=False,
+                 nms_class_agnostic=False):
         super(YOLOv2, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -31,6 +32,7 @@ class YOLOv2(nn.Module):
         self.topk = topk                               # topk
         self.stride = 32                               # 网络的最大步长
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         # ------------------- Anchor box -------------------
         self.anchor_size = torch.as_tensor(cfg['anchor_size']).float().view(-1, 2) # [A, 2]
         self.num_anchors = self.anchor_size.shape[0]
@@ -143,7 +145,7 @@ class YOLOv2(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels
 

models/detectors/yolov3/build.py

@@ -25,7 +25,8 @@ def build_yolov3(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         conf_thresh=args.conf_thresh,
         nms_thresh=args.nms_thresh,
         topk=args.topk,
-        deploy = deploy
+        deploy = deploy,
+        nms_class_agnostic = args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov3/yolov3.py

@@ -19,7 +19,8 @@ class YOLOv3(nn.Module):
                  topk=100,
                  nms_thresh=0.5,
                  trainable=False,
-                 deploy=False):
+                 deploy=False,
+                 nms_class_agnostic=False):
         super(YOLOv3, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -31,6 +32,7 @@ class YOLOv3(nn.Module):
         self.topk = topk                               # topk
         self.stride = [8, 16, 32]                      # 网络的输出步长
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         # ------------------- Anchor box -------------------
         self.num_levels = 3
         self.num_anchors = len(cfg['anchor_size']) // self.num_levels
@@ -147,7 +149,7 @@ class YOLOv3(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels
 

models/detectors/yolov4/build.py

@@ -25,7 +25,8 @@ def build_yolov4(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         conf_thresh=args.conf_thresh,
         nms_thresh=args.nms_thresh,
         topk=args.topk,
-        deploy = deploy
+        deploy=deploy,
+        nms_class_agnostic=args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov4/yolov4.py

@@ -19,7 +19,8 @@ class YOLOv4(nn.Module):
                  nms_thresh=0.5,
                  topk=100,
                  trainable=False,
-                 deploy=False):
+                 deploy=False,
+                 nms_class_agnostic=False):
         super(YOLOv4, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -31,6 +32,7 @@ class YOLOv4(nn.Module):
         self.topk = topk                               # topk
         self.stride = [8, 16, 32]                      # 网络的输出步长
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         # ------------------- Anchor box -------------------
         self.num_levels = 3
         self.num_anchors = len(cfg['anchor_size']) // self.num_levels

models/detectors/yolov5/build.py

@@ -25,7 +25,8 @@ def build_yolov5(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         nms_thresh=args.nms_thresh,
         topk=args.topk,
         trainable = trainable,
-        deploy = deploy
+        deploy=deploy,
+        nms_class_agnostic = args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov5/yolov5.py

@@ -17,7 +17,8 @@ class YOLOv5(nn.Module):
                  nms_thresh = 0.6,
                  trainable = False, 
                  topk = 1000,
-                 deploy = False):
+                 deploy = False,
+                 nms_class_agnostic = False):
         super(YOLOv5, self).__init__()
         # ---------------------- Basic Parameters ----------------------
         self.cfg = cfg
@@ -29,6 +30,7 @@ class YOLOv5(nn.Module):
         self.nms_thresh = nms_thresh
         self.topk = topk
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         
         # ------------------- Anchor box -------------------
         self.num_levels = 3
@@ -138,7 +140,7 @@ class YOLOv5(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels
 

models/detectors/yolov7/build.py

@@ -25,7 +25,8 @@ def build_yolov7(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
         nms_thresh = args.nms_thresh,
         topk = args.topk,
         trainable = trainable,
-        deploy = deploy
+        deploy = deploy,
+        nms_class_agnostic = args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolov7/yolov7.py

@@ -19,7 +19,8 @@ class YOLOv7(nn.Module):
                  topk=100,
                  nms_thresh=0.5,
                  trainable=False,
-                 deploy = False):
+                 deploy = False,
+                 nms_class_agnostic = False):
         super(YOLOv7, self).__init__()
         # ------------------- Basic parameters -------------------
         self.cfg = cfg                                 # 模型配置文件
@@ -31,6 +32,7 @@ class YOLOv7(nn.Module):
         self.topk = topk                               # topk
         self.stride = [8, 16, 32]                      # 网络的输出步长        
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
         # ------------------- Network Structure -------------------
         ## 主干网络
         self.backbone, feats_dim = build_backbone(cfg, trainable&cfg['pretrained'])
@@ -131,7 +133,7 @@ class YOLOv7(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels
 

models/detectors/yolox/build.py

@@ -25,7 +25,8 @@ def build_yolox(args, cfg, device, num_classes=80, trainable=False, deploy=False
         nms_thresh=args.nms_thresh,
         topk=args.topk,
         trainable = trainable,
-        deploy = deploy
+        deploy = deploy,
+        nms_class_agnostic = args.nms_class_agnostic
         )
 
     # -------------- Initialize YOLO --------------

models/detectors/yolox/yolox.py

@@ -17,7 +17,8 @@ class YOLOX(nn.Module):
                  nms_thresh = 0.6,
                  trainable = False, 
                  topk = 1000,
-                 deploy = False):
+                 deploy = False,
+                 nms_class_agnostic = False):
         super(YOLOX, self).__init__()
         # ---------------------- Basic Parameters ----------------------
         self.cfg = cfg
@@ -29,6 +30,7 @@ class YOLOX(nn.Module):
         self.nms_thresh = nms_thresh
         self.topk = topk
         self.deploy = deploy
+        self.nms_class_agnostic = nms_class_agnostic
                 
         # ------------------- Network Structure -------------------
         ## 主干网络
@@ -126,7 +128,7 @@ class YOLOX(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
 
         return bboxes, scores, labels