modify init

demo.py

@@ -60,8 +60,6 @@ def parse_args():
                         help='NMS threshold')
     parser.add_argument('--topk', default=1000, type=int,
                         help='topk candidates dets of each level before NMS')
-    parser.add_argument('--max_dets', default=300, type=int,
-                        help='max number of dets after NMS')
     parser.add_argument("--deploy", action="store_true", default=False,
                         help="deploy mode or not")
     parser.add_argument('--fuse_repconv', action='store_true', default=False,

eval.py

@@ -38,8 +38,6 @@ def parse_args():
                         help='NMS threshold')
     parser.add_argument('--topk', default=1000, type=int,
                         help='topk candidates dets of each level before NMS')
-    parser.add_argument('--max_dets', default=300, type=int,
-                        help='max number of dets after NMS')
     parser.add_argument("--no_decode", action="store_true", default=False,
                         help="not decode in inference or yes")
     parser.add_argument('--fuse_conv_bn', action='store_true', default=False,

models/detectors/yolov8/build.py

@@ -24,7 +24,6 @@ def build_yolov8(args, cfg, device, num_classes=80, trainable=False, deploy=Fals
                    conf_thresh        = args.conf_thresh,
                    nms_thresh         = args.nms_thresh,
                    topk               = args.topk,
-                   max_dets           = args.max_dets,
                    deploy             = deploy,
                    nms_class_agnostic = args.nms_class_agnostic
                    )

models/detectors/yolov8/yolov8.py

@@ -22,7 +22,6 @@ class YOLOv8(nn.Module):
                  conf_thresh = 0.01,
                  nms_thresh  = 0.5,
                  topk        = 1000,
-                 max_dets    = 300,
                  trainable   = False,
                  deploy      = False,
                  nms_class_agnostic = False):
@@ -39,7 +38,6 @@ class YOLOv8(nn.Module):
         self.num_levels = len(self.strides)
         self.num_classes = num_classes
         self.topk = topk
-        self.max_dets = max_dets
         self.deploy = deploy
         self.nms_class_agnostic = nms_class_agnostic
         
@@ -118,7 +116,7 @@ class YOLOv8(nn.Module):
 
         # nms
         scores, labels, bboxes = multiclass_nms(
-            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic, self.max_dets)
+            scores, labels, bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)
         
         return bboxes, scores, labels
 

models/detectors/yolov8/yolov8_head.py

@@ -59,6 +59,15 @@ class SingleLevelHead(nn.Module):
         self.cls_feats = nn.Sequential(*cls_feats)
         self.reg_feats = nn.Sequential(*reg_feats)
 
+        self.init_weights()
+        
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
 
     def forward(self, x):
         """

models/detectors/yolov8/yolov8_pafpn.py

@@ -68,6 +68,16 @@ class Yolov8PaFPN(nn.Module):
                                                   depthwise    = depthwise,
                                                   )
 
+        self.init_weights()
+        
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, features):
         c3, c4, c5 = features
 

models/detectors/yolov8/yolov8_pred.py

@@ -31,9 +31,6 @@ class SingleLevelPredLayer(nn.Module):
         b = self.reg_pred.bias.view(-1, )
         b.data.fill_(1.0)
         self.reg_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
-        w = self.reg_pred.weight
-        w.data.fill_(0.)
-        self.reg_pred.weight = torch.nn.Parameter(w, requires_grad=True)
 
     def forward(self, cls_feat, reg_feat):
         """

test.py

@@ -47,10 +47,8 @@ def parse_args():
                         help='confidence threshold')
     parser.add_argument('-nt', '--nms_thresh', default=0.5, type=float,
                         help='NMS threshold')
-    parser.add_argument('--topk', default=1000, type=int,
+    parser.add_argument('--topk', default=100, type=int,
                         help='topk candidates dets of each level before NMS')
-    parser.add_argument('--max_dets', default=300, type=int,
-                        help='max number of dets after NMS')
     parser.add_argument("--no_decode", action="store_true", default=False,
                         help="not decode in inference or yes")
     parser.add_argument('--fuse_conv_bn', action='store_true', default=False,

train.py

@@ -67,8 +67,6 @@ def parse_args():
                         help='NMS threshold')
     parser.add_argument('--topk', default=1000, type=int,
                         help='topk candidates dets of each level before NMS')
-    parser.add_argument('--max_dets', default=300, type=int,
-                        help='max number of dets after NMS')
     parser.add_argument('-p', '--pretrained', default=None, type=str,
                         help='load pretrained weight')
     parser.add_argument('-r', '--resume', default=None, type=str,

utils/misc.py

@@ -290,22 +290,17 @@ def nms(bboxes, scores, nms_thresh):
     return keep
 
 ## class-agnostic NMS 
-def multiclass_nms_class_agnostic(scores, labels, bboxes, nms_thresh, max_dets=300):
+def multiclass_nms_class_agnostic(scores, labels, bboxes, nms_thresh):
     # nms
     keep = nms(bboxes, scores, nms_thresh)
     scores = scores[keep]
     labels = labels[keep]
     bboxes = bboxes[keep]
 
-    # max dets
-    scores = scores[:max_dets]
-    labels = labels[:max_dets]
-    bboxes = bboxes[:max_dets]
-
     return scores, labels, bboxes
 
 ## class-aware NMS 
-def multiclass_nms_class_aware(scores, labels, bboxes, nms_thresh, num_classes, max_dets=300):
+def multiclass_nms_class_aware(scores, labels, bboxes, nms_thresh, num_classes):
     # nms
     keep = np.zeros(len(bboxes), dtype=np.int32)
     for i in range(num_classes):
@@ -321,19 +316,14 @@ def multiclass_nms_class_aware(scores, labels, bboxes, nms_thresh, num_classes,
     labels = labels[keep]
     bboxes = bboxes[keep]
 
-    # max dets
-    scores = scores[:max_dets]
-    labels = labels[:max_dets]
-    bboxes = bboxes[:max_dets]
-
     return scores, labels, bboxes
 
 ## multi-class NMS 
-def multiclass_nms(scores, labels, bboxes, nms_thresh, num_classes, class_agnostic=False, max_dets=300):
+def multiclass_nms(scores, labels, bboxes, nms_thresh, num_classes, class_agnostic=False):
     if class_agnostic:
-        return multiclass_nms_class_agnostic(scores, labels, bboxes, nms_thresh, max_dets)
+        return multiclass_nms_class_agnostic(scores, labels, bboxes, nms_thresh)
     else:
-        return multiclass_nms_class_aware(scores, labels, bboxes, nms_thresh, num_classes, max_dets)
+        return multiclass_nms_class_aware(scores, labels, bboxes, nms_thresh, num_classes)
 
 
 # ---------------------------- Processor for Deployment ----------------------------