add YOLOv7-Plus

config/__init__.py

@@ -5,7 +5,7 @@ from .yolov3_config import yolov3_cfg
 from .yolov4_config import yolov4_cfg
 from .yolov5_config import yolov5_cfg
 from .yolov7_config import yolov7_cfg
-from .yolov8_config import yolov8_cfg
+from .yolov7_plus_config import yolov7_plus_cfg
 from .yolox_config import yolox_cfg
 
 
@@ -31,8 +31,8 @@ def build_model_config(args):
     elif args.model in ['yolov7_t', 'yolov7_l', 'yolov7_x']:
         cfg = yolov7_cfg[args.model]
     # YOLOv8
-    elif args.model in ['yolov8_n', 'yolov8_s', 'yolov8_m', 'yolov8_l', 'yolov8_x']:
-        cfg = yolov8_cfg[args.model]
+    elif args.model in ['yolov7_plus_n', 'yolov7_plus_s', 'yolov7_plus_m', 'yolov7_plus_l', 'yolov7_plus_x']:
+        cfg = yolov7_plus_cfg[args.model]
     # YOLOX
     elif args.model == 'yolox':
         cfg = yolox_cfg

config/yolov8_config.py → config/yolov7_plus_config.py

@@ -1,11 +1,12 @@
 # yolov8 config
 
-yolov8_cfg = {
-    'yolov8_n':{
+yolov7_plus_cfg = {
+    'yolov7_plus_n':{
         # input
         'trans_type': 'yolov5_tiny',
-        'multi_scale': [0.5, 1.5],   # 320 -> 960
-        # model
+        'multi_scale': [0.5, 1.0],   # 320 -> 640
+        # ----------------- Model config -----------------
+        ## Backbone
         'backbone': 'elan_cspnet',
         'pretrained': True,
         'bk_act': 'silu',
@@ -15,58 +16,69 @@ yolov8_cfg = {
         'depth': 0.34,
         'ratio': 2.0,
         'stride': [8, 16, 32],  # P3, P4, P5
-        # neck
+        ## Neck: SPP
         'neck': 'sppf',
         'expand_ratio': 0.5,
         'pooling_size': 5,
         'neck_act': 'silu',
         'neck_norm': 'BN',
         'neck_depthwise': False,
-        # fpn
-        'fpn': 'yolov8_pafpn',
+        ## Neck: FPN
+        'fpn': 'yolov7_plus_pafpn',
+        'fpn_reduce_layer': 'Conv',
+        'fpn_downsample_layer': 'Conv',
+        'fpn_core_block': 'ELAN_CSPBlock',
         'fpn_act': 'silu',
         'fpn_norm': 'BN',
         'fpn_depthwise': False,
-        # head
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        ## Head
         'head': 'decoupled_head',
         'head_act': 'silu',
         'head_norm': 'BN',
         'num_cls_head': 2,
         'num_reg_head': 2,
         'head_depthwise': False,
-        'reg_max': 16,
-        # matcher
-        'matcher': {'topk': 10,
-                    'alpha': 0.5,
-                    'beta': 6.0},
-        # loss weight
-        'cls_loss': 'bce', # vfl (optional)
+        # ----------------- Label Assignment config -----------------
+        'matcher': {
+            ## For fixed assigner
+            'anchor_thresh': 4.0,
+            ## For dynamic assigner
+            'topk': 10,
+            'alpha': 0.5,
+            'beta': 6.0},
+        # ----------------- Loss config -----------------
+        'cls_loss': 'bce',
         'loss_cls_weight': 0.5,
         'loss_iou_weight': 7.5,
-        'loss_dfl_weight': 1.5,
-        # training configuration
+        # ----------------- Train config -----------------
+        ## stop strong augment
         'no_aug_epoch': 20,
-        # optimizer
+        ## optimizer
         'optimizer': 'sgd',        # optional: sgd, adamw
         'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
         'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
         'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-        # model EMA
+        ## Model EMA
         'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
         'ema_tau': 2000,
-        # lr schedule
+        ## LR schedule
         'scheduler': 'linear',
-        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.004
+        'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
         'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
+        ## WarmUpLR schedule
         'warmup_momentum': 0.8,
         'warmup_bias_lr': 0.1,
     },
 
-    'yolov8_s':{
+    'yolov7_plus_s':{
         # input
         'trans_type': 'yolov5_small',
-        'multi_scale': [0.5, 1.5],   # 320 -> 960
-        # model
+        'multi_scale': [0.5, 1.0],   # 320 -> 640
+        # ----------------- Model config 
+        # Backbone
         'backbone': 'elan_cspnet',
         'pretrained': True,
         'bk_act': 'silu',
@@ -76,58 +88,69 @@ yolov8_cfg = {
         'depth': 0.34,
         'ratio': 2.0,
         'stride': [8, 16, 32],  # P3, P4, P5
-        # neck
+        # Neck: SPP
         'neck': 'sppf',
         'expand_ratio': 0.5,
         'pooling_size': 5,
         'neck_act': 'silu',
         'neck_norm': 'BN',
         'neck_depthwise': False,
-        # fpn
-        'fpn': 'yolov8_pafpn',
+        # Neck: FPN
+        'fpn': 'yolov7_plus_pafpn',
+        'fpn_reduce_layer': 'Conv',
+        'fpn_downsample_layer': 'Conv',
+        'fpn_core_block': 'ELAN_CSPBlock',
         'fpn_act': 'silu',
         'fpn_norm': 'BN',
         'fpn_depthwise': False,
-        # head
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # Head
         'head': 'decoupled_head',
         'head_act': 'silu',
         'head_norm': 'BN',
         'num_cls_head': 2,
         'num_reg_head': 2,
         'head_depthwise': False,
-        'reg_max': 16,
-        # matcher
-        'matcher': {'topk': 10,
-                    'alpha': 0.5,
-                    'beta': 6.0},
-        # loss weight
-        'cls_loss': 'bce', # vfl (optional)
+        # ----------------- Label Assignment config -----------------
+        'matcher': {
+            ## For fixed assigner
+            'anchor_thresh': 4.0,
+            ## For dynamic assigner
+            'topk': 10,
+            'alpha': 0.5,
+            'beta': 6.0},
+        # ----------------- Loss config -----------------
+        'cls_loss': 'bce',
         'loss_cls_weight': 0.5,
         'loss_iou_weight': 7.5,
-        'loss_dfl_weight': 1.5,
-        # training configuration
+        # ----------------- Train config -----------------
+        # stop strong augment
         'no_aug_epoch': 20,
-        # optimizer
+        ## optimizer
         'optimizer': 'sgd',        # optional: sgd, adamw
         'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
         'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
         'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-        # model EMA
+        ## Model EMA
         'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
         'ema_tau': 2000,
-        # lr schedule
+        ## LR schedule
         'scheduler': 'linear',
-        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.004
+        'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
         'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
+        ## WarmUpLR schedule
         'warmup_momentum': 0.8,
         'warmup_bias_lr': 0.1,
     },
 
-    'yolov8_m':{
+    'yolov7_plus_m':{
         # input
         'trans_type': 'yolov5_medium',
-        'multi_scale': [0.5, 1.5],   # 320 -> 960
-        # model
+        'multi_scale': [0.5, 1.0],   # 320 -> 640
+        # ----------------- Model config 
+        # Backbone
         'backbone': 'elan_cspnet',
         'pretrained': True,
         'bk_act': 'silu',
@@ -137,58 +160,69 @@ yolov8_cfg = {
         'depth': 0.67,
         'ratio': 1.5,
         'stride': [8, 16, 32],  # P3, P4, P5
-        # neck
+        # Neck: SPP
         'neck': 'sppf',
         'expand_ratio': 0.5,
         'pooling_size': 5,
         'neck_act': 'silu',
         'neck_norm': 'BN',
         'neck_depthwise': False,
-        # fpn
-        'fpn': 'yolov8_pafpn',
+        # Neck: FPN
+        'fpn': 'yolov7_plus_pafpn',
+        'fpn_reduce_layer': 'Conv',
+        'fpn_downsample_layer': 'Conv',
+        'fpn_core_block': 'ELAN_CSPBlock',
         'fpn_act': 'silu',
         'fpn_norm': 'BN',
         'fpn_depthwise': False,
-        # head
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # Head
         'head': 'decoupled_head',
         'head_act': 'silu',
         'head_norm': 'BN',
         'num_cls_head': 2,
         'num_reg_head': 2,
         'head_depthwise': False,
-        'reg_max': 16,
-        # matcher
-        'matcher': {'topk': 10,
-                    'alpha': 0.5,
-                    'beta': 6.0},
-        # loss weight
-        'cls_loss': 'bce', # vfl (optional)
+        # ----------------- Label Assignment config -----------------
+        'matcher': {
+            ## For fixed assigner
+            'anchor_thresh': 4.0,
+            ## For dynamic assigner
+            'topk': 10,
+            'alpha': 0.5,
+            'beta': 6.0},
+        # ----------------- Loss config -----------------
+        'cls_loss': 'bce',
         'loss_cls_weight': 0.5,
         'loss_iou_weight': 7.5,
-        'loss_dfl_weight': 1.5,
-        # training configuration
+        # ----------------- Train config -----------------
+        # stop strong augment
         'no_aug_epoch': 20,
-        # optimizer
+        ## optimizer
         'optimizer': 'sgd',        # optional: sgd, adamw
         'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
         'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
         'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-        # model EMA
+        ## Model EMA
         'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
         'ema_tau': 2000,
-        # lr schedule
+        ## LR schedule
         'scheduler': 'linear',
-        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.004
+        'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
         'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
+        ## WarmUpLR schedule
         'warmup_momentum': 0.8,
         'warmup_bias_lr': 0.1,
     },
 
-    'yolov8_l':{
+    'yolov7_plus_l':{
         # input
         'trans_type': 'yolov5_large',
-        'multi_scale': [0.5, 1.5],   # 320 -> 960
-        # model
+        'multi_scale': [0.5, 1.0],   # 320 -> 640
+        # ----------------- Model config 
+        # Backbone
         'backbone': 'elan_cspnet',
         'pretrained': True,
         'bk_act': 'silu',
@@ -198,58 +232,69 @@ yolov8_cfg = {
         'depth': 1.0,
         'ratio': 1.0,
         'stride': [8, 16, 32],  # P3, P4, P5
-        # neck
+        # Neck: SPP
         'neck': 'sppf',
         'expand_ratio': 0.5,
         'pooling_size': 5,
         'neck_act': 'silu',
         'neck_norm': 'BN',
         'neck_depthwise': False,
-        # fpn
-        'fpn': 'yolov8_pafpn',
+        # Neck: FPN
+        'fpn': 'yolov7_plus_pafpn',
+        'fpn_reduce_layer': 'Conv',
+        'fpn_downsample_layer': 'Conv',
+        'fpn_core_block': 'ELAN_CSPBlock',
         'fpn_act': 'silu',
         'fpn_norm': 'BN',
         'fpn_depthwise': False,
-        # head
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # Head
         'head': 'decoupled_head',
         'head_act': 'silu',
         'head_norm': 'BN',
         'num_cls_head': 2,
         'num_reg_head': 2,
         'head_depthwise': False,
-        'reg_max': 16,
-        # matcher
-        'matcher': {'topk': 10,
-                    'alpha': 0.5,
-                    'beta': 6.0},
-        # loss weight
-        'cls_loss': 'bce', # vfl (optional)
+        # ----------------- Label Assignment config -----------------
+        'matcher': {
+            ## For fixed assigner
+            'anchor_thresh': 4.0,
+            ## For dynamic assigner
+            'topk': 10,
+            'alpha': 0.5,
+            'beta': 6.0},
+        # ----------------- Loss config -----------------
+        'cls_loss': 'bce',
         'loss_cls_weight': 0.5,
         'loss_iou_weight': 7.5,
-        'loss_dfl_weight': 1.5,
-        # training configuration
+        # ----------------- Train config -----------------
+        # stop strong augment
         'no_aug_epoch': 20,
-        # optimizer
+        ## optimizer
         'optimizer': 'sgd',        # optional: sgd, adamw
         'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
         'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
         'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-        # model EMA
+        ## Model EMA
         'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
         'ema_tau': 2000,
-        # lr schedule
+        ## LR schedule
         'scheduler': 'linear',
-        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.004
+        'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
         'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
+        ## WarmUpLR schedule
         'warmup_momentum': 0.8,
         'warmup_bias_lr': 0.1,
     },
 
-    'yolov8_x':{
+    'yolov7_plus_x':{
         # input
         'trans_type': 'yolov5_huge',
-        'multi_scale': [0.5, 1.5],   # 320 -> 960
-        # model
+        'multi_scale': [0.5, 1.0],   # 320 -> 640
+        # ----------------- Model config 
+        # Backbone
         'backbone': 'elan_cspnet',
         'pretrained': True,
         'bk_act': 'silu',
@@ -259,49 +304,59 @@ yolov8_cfg = {
         'depth': 1.0,
         'ratio': 1.0,
         'stride': [8, 16, 32],  # P3, P4, P5
-        # neck
+        # Neck: SPP
         'neck': 'sppf',
         'expand_ratio': 0.5,
         'pooling_size': 5,
         'neck_act': 'silu',
         'neck_norm': 'BN',
         'neck_depthwise': False,
-        # fpn
-        'fpn': 'yolov8_pafpn',
+        # Neck: FPN
+        'fpn': 'yolov7_plus_pafpn',
+        'fpn_reduce_layer': 'Conv',
+        'fpn_downsample_layer': 'Conv',
+        'fpn_core_block': 'ELAN_CSPBlock',
         'fpn_act': 'silu',
         'fpn_norm': 'BN',
         'fpn_depthwise': False,
-        # head
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # Head
         'head': 'decoupled_head',
         'head_act': 'silu',
         'head_norm': 'BN',
         'num_cls_head': 2,
         'num_reg_head': 2,
         'head_depthwise': False,
-        'reg_max': 16,
-        # matcher
-        'matcher': {'topk': 10,
-                    'alpha': 0.5,
-                    'beta': 6.0},
-        # loss weight
-        'cls_loss': 'bce', # vfl (optional)
+        # ----------------- Label Assignment config -----------------
+        'matcher': {
+            ## For fixed assigner
+            'anchor_thresh': 4.0,
+            ## For dynamic assigner
+            'topk': 10,
+            'alpha': 0.5,
+            'beta': 6.0},
+        # ----------------- Loss config -----------------
+        'cls_loss': 'bce',
         'loss_cls_weight': 0.5,
         'loss_iou_weight': 7.5,
-        'loss_dfl_weight': 1.5,
-        # training configuration
+        # ----------------- Train config -----------------
+        # stop strong augment
         'no_aug_epoch': 20,
-        # optimizer
+        ## optimizer
         'optimizer': 'sgd',        # optional: sgd, adamw
         'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
         'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
         'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-        # model EMA
+        ## Model EMA
         'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
         'ema_tau': 2000,
-        # lr schedule
+        ## LR schedule
         'scheduler': 'linear',
-        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.004
+        'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
         'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
+        ## WarmUpLR schedule
         'warmup_momentum': 0.8,
         'warmup_bias_lr': 0.1,
     },

engine.py

@@ -11,6 +11,22 @@ from utils import distributed_utils
 from utils.vis_tools import vis_data
 
 
+def refine_targets(targets, min_box_size):
+    # rescale targets
+    for tgt in targets:
+        boxes = tgt["boxes"].clone()
+        labels = tgt["labels"].clone()
+        # refine tgt
+        tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
+        min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
+        keep = (min_tgt_size >= min_box_size)
+
+        tgt["boxes"] = boxes[keep]
+        tgt["labels"] = labels[keep]
+    
+    return targets
+
+
 def rescale_image_targets(images, targets, stride, min_box_size, multi_scale_range=[0.5, 1.5]):
     """
         Deployed for Multi scale trick.
@@ -84,10 +100,6 @@ def train_one_epoch(epoch,
                 if 'momentum' in x:
                     x['momentum'] = np.interp(ni, xi, [cfg['warmup_momentum'], cfg['momentum']])
                             
-        # visualize train targets
-        if args.vis_tgt:
-            vis_data(images, targets)
-
         # to device
         images = images.to(device, non_blocking=True).float() / 255.
 
@@ -95,12 +107,18 @@ def train_one_epoch(epoch,
         if args.multi_scale:
             images, targets, img_size = rescale_image_targets(
                 images, targets, model.stride, args.min_box_size, cfg['multi_scale'])
+        else:
+            targets = refine_targets(targets, args.min_box_size)
             
+        # visualize train targets
+        if args.vis_tgt:
+            vis_data(images*255, targets)
+
         # inference
         with torch.cuda.amp.autocast(enabled=args.fp16):
             outputs = model(images)
             # loss
-            loss_dict = criterion(outputs=outputs, targets=targets)
+            loss_dict = criterion(outputs, targets, epoch)
             losses = loss_dict['losses']
             losses *= images.shape[0]  # loss * bs
 

models/detectors/__init__.py

@@ -8,7 +8,7 @@ from .yolov3.build import build_yolov3
 from .yolov4.build import build_yolov4
 from .yolov5.build import build_yolov5
 from .yolov7.build import build_yolov7
-from .yolov8.build import build_yolov8
+from .yolov7_plus.build import build_yolov7_plus
 from .yolox.build import build_yolox
 
 
@@ -42,9 +42,9 @@ def build_model(args,
     elif args.model in ['yolov7_t', 'yolov7_l', 'yolov7_x']:
         model, criterion = build_yolov7(
             args, model_cfg, device, num_classes, trainable)
-    # YOLOv8
-    elif args.model in ['yolov8_n', 'yolov8_s', 'yolov8_m', 'yolov8_l', 'yolov8_x']:
-        model, criterion = build_yolov8(
+    # YOLOv7-Plus
+    elif args.model in ['yolov7_plus_n', 'yolov7_plus_s', 'yolov7_plus_m', 'yolov7_plus_l', 'yolov7_plus_x']:
+        model, criterion = build_yolov7_plus(
             args, model_cfg, device, num_classes, trainable)
     # YOLOX   
     elif args.model == 'yolox':

models/detectors/yolov1/loss.py

@@ -41,7 +41,7 @@ class Criterion(object):
         return loss_box
 
 
-    def __call__(self, outputs, targets):
+    def __call__(self, outputs, targets, epoch=0):
         device = outputs['pred_cls'][0].device
         stride = outputs['stride']
         fmp_size = outputs['fmp_size']

models/detectors/yolov2/loss.py

@@ -42,7 +42,7 @@ class Criterion(object):
         return loss_box, ious
 
 
-    def __call__(self, outputs, targets):
+    def __call__(self, outputs, targets, epoch=0):
         device = outputs['pred_cls'].device
         stride = outputs['stride']
         fmp_size = outputs['fmp_size']

models/detectors/yolov3/loss.py

@@ -42,7 +42,7 @@ class Criterion(object):
         return loss_box, ious
 
 
-    def __call__(self, outputs, targets):
+    def __call__(self, outputs, targets, epoch=0):
         device = outputs['pred_cls'][0].device
         fpn_strides = outputs['strides']
         fmp_sizes = outputs['fmp_sizes']

models/detectors/yolov4/loss.py

@@ -42,7 +42,7 @@ class Criterion(object):
         return loss_box, ious
 
 
-    def __call__(self, outputs, targets):
+    def __call__(self, outputs, targets, epoch=0):
         device = outputs['pred_cls'][0].device
         fpn_strides = outputs['strides']
         fmp_sizes = outputs['fmp_sizes']

models/detectors/yolov5/loss.py

@@ -42,7 +42,7 @@ class Criterion(object):
         return loss_box, ious
 
 
-    def __call__(self, outputs, targets):
+    def __call__(self, outputs, targets, epoch=0):
         device = outputs['pred_cls'][0].device
         fpn_strides = outputs['strides']
         fmp_sizes = outputs['fmp_sizes']

models/detectors/yolov7/loss.py

@@ -51,7 +51,7 @@ class Criterion(object):
         return loss_box
 
 
-    def __call__(self, outputs, targets):        
+    def __call__(self, outputs, targets, epoch=0):        
         """
             outputs['pred_obj']: List(Tensor) [B, M, 1]
             outputs['pred_cls']: List(Tensor) [B, M, C]

models/detectors/yolov8/build.py → models/detectors/yolov7_plus/build.py

@@ -5,11 +5,11 @@ import torch
 import torch.nn as nn
 
 from .loss import build_criterion
-from .yolov8 import YOLOv8
+from .yolov7_plus import YOLOv7_Plus
 
 
 # build object detector
-def build_yolov8(args, cfg, device, num_classes=80, trainable=False):
+def build_yolov7_plus(args, cfg, device, num_classes=80, trainable=False):
     print('==============================')
     print('Build {} ...'.format(args.model.upper()))
     
@@ -17,7 +17,7 @@ def build_yolov8(args, cfg, device, num_classes=80, trainable=False):
     print('Model Configuration: \n', cfg)
     
     # -------------- Build YOLO --------------
-    model = YOLOv8(
+    model = YOLOv7_Plus(
         cfg=cfg,
         device=device, 
         num_classes=num_classes,
@@ -54,5 +54,5 @@ def build_yolov8(args, cfg, device, num_classes=80, trainable=False):
     criterion = None
     if trainable:
         # build criterion for training
-        criterion = build_criterion(cfg, device, num_classes)
+        criterion = build_criterion(cfg, device, num_classes, warmup_epoch=args.wp_epoch)
     return model, criterion

models/detectors/yolov8/loss.py → models/detectors/yolov7_plus/loss.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from .matcher import TaskAlignedAssigner
-from utils.box_ops import bbox2dist, bbox_iou
+from .matcher import TaskAlignedAssigner, Yolov5Matcher
+from utils.box_ops import bbox_iou, get_ious
+from utils.distributed_utils import get_world_size, is_dist_avail_and_initialized
 
 
 
@@ -10,22 +11,32 @@ class Criterion(object):
     def __init__(self, 
                  cfg, 
                  device, 
-                 num_classes=80):
+                 num_classes=80,
+                 warmup_epoch=1):
+        # ------------------ Basic Parameters ------------------
         self.cfg = cfg
         self.device = device
         self.num_classes = num_classes
-        self.reg_max = cfg['reg_max']
-        self.use_dfl = cfg['reg_max'] > 1
-        # loss
+        self.warmup_epoch = warmup_epoch
+        self.warmup_stage = True
+        # ------------------ Loss Parameters ------------------
+        ## loss function
         self.cls_lossf = ClassificationLoss(cfg, reduction='none')
-        self.reg_lossf = RegressionLoss(num_classes, cfg['reg_max'] - 1, self.use_dfl)
-        # loss weight
+        self.reg_lossf = RegressionLoss(num_classes)
+        ## loss coeff
         self.loss_cls_weight = cfg['loss_cls_weight']
         self.loss_iou_weight = cfg['loss_iou_weight']
-        self.loss_dfl_weight = cfg['loss_dfl_weight']
-        # matcher
+        # ------------------ Label Assigner ------------------
         matcher_config = cfg['matcher']
-        self.matcher = TaskAlignedAssigner(
+        ## matcher-1
+        self.fixed_matcher = Yolov5Matcher(
+            num_classes=num_classes, 
+            num_anchors=3, 
+            anchor_size=cfg['anchor_size'],
+            anchor_theshold=matcher_config['anchor_thresh']
+            )
+        ## matcher-2
+        self.dynamic_matcher = TaskAlignedAssigner(
             topk=matcher_config['topk'],
             num_classes=num_classes,
             alpha=matcher_config['alpha'],
@@ -33,28 +44,67 @@ class Criterion(object):
             )
 
 
-    def __call__(self, outputs, targets):        
-        """
-            outputs['pred_cls']: List(Tensor) [B, M, C]
-            outputs['pred_regs']: List(Tensor) [B, M, 4*(reg_max+1)]
-            outputs['pred_boxs']: List(Tensor) [B, M, 4]
-            outputs['anchors']: List(Tensor) [M, 2]
-            outputs['strides']: List(Int) [8, 16, 32] output stride
-            outputs['stride_tensor']: List(Tensor) [M, 1]
-            targets: (List) [dict{'boxes': [...], 
-                                 'labels': [...], 
-                                 'orig_size': ...}, ...]
-        """
+    def fixed_assignment_loss(self, outputs, targets):
+        device = outputs['pred_cls'][0].device
+        fpn_strides = outputs['strides']
+        fmp_sizes = outputs['fmp_sizes']
+        (
+            gt_objectness, 
+            gt_classes, 
+            gt_bboxes,
+            ) = self.fixed_matcher(fmp_sizes=fmp_sizes, 
+                                   fpn_strides=fpn_strides, 
+                                   targets=targets)
+        # List[B, M, C] -> [B, M, C] -> [BM, C]
+        pred_cls = torch.cat(outputs['pred_cls'], dim=1).view(-1, self.num_classes)    # [BM, C]
+        pred_box = torch.cat(outputs['pred_box'], dim=1).view(-1, 4)                   # [BM, 4]
+       
+        gt_objectness = gt_objectness.view(-1).to(device).float()               # [BM,]
+        gt_classes = gt_classes.view(-1, self.num_classes).to(device).float()   # [BM, C]
+        gt_bboxes = gt_bboxes.view(-1, 4).to(device).float()                    # [BM, 4]
+
+        pos_masks = (gt_objectness > 0)
+        num_fgs = pos_masks.sum()
+
+        if is_dist_avail_and_initialized():
+            torch.distributed.all_reduce(num_fgs)
+        num_fgs = (num_fgs / get_world_size()).clamp(1.0)
+
+        # box loss
+        ious = get_ious(pred_box[pos_masks],
+                        gt_bboxes[pos_masks],
+                        box_mode="xyxy",
+                        iou_type='giou')
+        loss_box = 1.0 - ious
+        loss_box = loss_box.sum() / num_fgs
+        
+        # cls loss
+        gt_classes[pos_masks] = gt_classes[pos_masks] * ious.unsqueeze(-1).clamp(0.)
+        loss_cls = F.binary_cross_entropy_with_logits(pred_cls, gt_classes, reduction='none')
+        loss_cls = loss_cls.sum() / num_fgs
+
+        # total loss
+        losses = self.loss_cls_weight * loss_cls + \
+                 self.loss_iou_weight * loss_box
+
+        loss_dict = dict(
+                loss_cls = loss_cls,
+                loss_box = loss_box,
+                losses = losses
+        )
+
+        return loss_dict
+
+
+    def dynamic_assignment_loss(self, outputs, targets):
         bs = outputs['pred_cls'][0].shape[0]
         device = outputs['pred_cls'][0].device
-        strides = outputs['stride_tensor']
         anchors = outputs['anchors']
         anchors = torch.cat(anchors, dim=0)
         num_anchors = anchors.shape[0]
 
         # preds: [B, M, C]
         cls_preds = torch.cat(outputs['pred_cls'], dim=1)
-        reg_preds = torch.cat(outputs['pred_reg'], dim=1)
         box_preds = torch.cat(outputs['pred_box'], dim=1)
         
         # label assignment
@@ -81,10 +131,10 @@ class Criterion(object):
                     gt_score,   #[1, M, C]
                     fg_mask,    #[1, M,]
                     _
-                ) = self.matcher(
+                ) = self.dynamic_matcher(
                     pd_scores = cls_preds[batch_idx:batch_idx+1].detach().sigmoid(), 
                     pd_bboxes = box_preds[batch_idx:batch_idx+1].detach(),
-                    anc_points = anchors,
+                    anc_points = anchors[..., :2],
                     gt_labels = tgt_labels,
                     gt_bboxes = tgt_boxs
                     )
@@ -102,46 +152,56 @@ class Criterion(object):
         loss_cls = self.cls_lossf(cls_preds, gt_score_targets)
 
         # reg loss
-        anchors = anchors[None].repeat(bs, 1, 1).view(-1, 2)                           # [BM, 2]
-        strides = torch.cat(strides, dim=0).unsqueeze(0).repeat(bs, 1, 1).view(-1, 1)  # [BM, 1]
         bbox_weight = gt_score_targets[fg_masks].sum(-1, keepdim=True)                 # [BM, 1]
-        reg_preds = reg_preds.view(-1, 4*self.reg_max)                                 # [BM, 4*(reg_max + 1)]
         box_preds = box_preds.view(-1, 4)                                              # [BM, 4]
-        loss_iou, loss_dfl = self.reg_lossf(
-            pred_regs = reg_preds,
+        loss_iou = self.reg_lossf(
             pred_boxs = box_preds,
-            anchors = anchors,
             gt_boxs = gt_bbox_targets,
             bbox_weight = bbox_weight,
-            fg_masks = fg_masks,
-            strides = strides,
+            fg_masks = fg_masks
             )
-        
+
+        num_fgs = gt_score_targets.sum()
+        if is_dist_avail_and_initialized():
+            torch.distributed.all_reduce(num_fgs)
+        num_fgs = (num_fgs / get_world_size()).clamp(1.0)
+
         # normalize loss
-        gt_score_targets_sum = max(gt_score_targets.sum(), 1)
-        loss_cls = loss_cls.sum() / gt_score_targets_sum
-        loss_iou = loss_iou.sum() / gt_score_targets_sum
-        loss_dfl = loss_dfl.sum() / gt_score_targets_sum
+        loss_cls = loss_cls.sum() / num_fgs
+        loss_iou = loss_iou.sum() / num_fgs
 
         # total loss
         losses = loss_cls * self.loss_cls_weight + \
                  loss_iou * self.loss_iou_weight
-        if self.use_dfl:
-            losses += loss_dfl * self.loss_dfl_weight
-            loss_dict = dict(
-                    loss_cls = loss_cls,
-                    loss_iou = loss_iou,
-                    loss_dfl = loss_dfl,
-                    losses = losses
-            )
-        else:
-            loss_dict = dict(
-                    loss_cls = loss_cls,
-                    loss_iou = loss_iou,
-                    losses = losses
-            )
+        loss_dict = dict(
+                loss_cls = loss_cls,
+                loss_iou = loss_iou,
+                losses = losses
+        )
 
         return loss_dict
+
+
+    def __call__(self, outputs, targets, epoch=0):        
+        """
+            outputs['pred_cls']: List(Tensor) [B, M, C]
+            outputs['pred_regs']: List(Tensor) [B, M, 4*(reg_max+1)]
+            outputs['pred_boxs']: List(Tensor) [B, M, 4]
+            outputs['anchors']: List(Tensor) [M, 2]
+            outputs['strides']: List(Int) [8, 16, 32] output stride
+            outputs['stride_tensor']: List(Tensor) [M, 1]
+            targets: (List) [dict{'boxes': [...], 
+                                 'labels': [...], 
+                                 'orig_size': ...}, ...]
+        """
+        if epoch >= self.warmup_epoch:
+            print('Switch to Dynamic Label Assignment.')
+            self.warmup_stage = False
+
+        if self.warmup_stage:
+            return self.fixed_assignment_loss(outputs, targets)
+        else:
+            return self.dynamic_assignment_loss(outputs, targets)
     
 
 class ClassificationLoss(nn.Module):
@@ -149,9 +209,6 @@ class ClassificationLoss(nn.Module):
         super(ClassificationLoss, self).__init__()
         self.cfg = cfg
         self.reduction = reduction
-        # For VFL
-        self.alpha = 0.75
-        self.gamma = 2.0
 
 
     def binary_cross_entropy(self, pred_logits, gt_score):
@@ -172,38 +229,14 @@ class ClassificationLoss(nn.Module):
 
 
 class RegressionLoss(nn.Module):
-    def __init__(self, num_classes, reg_max, use_dfl):
+    def __init__(self, num_classes):
         super(RegressionLoss, self).__init__()
         self.num_classes = num_classes
-        self.reg_max = reg_max
-        self.use_dfl = use_dfl
-
-
-    def df_loss(self, pred_regs, target):
-        gt_left = target.to(torch.long)
-        gt_right = gt_left + 1
-        weight_left = gt_right.to(torch.float) - target
-        weight_right = 1 - weight_left
-        # loss left
-        loss_left = F.cross_entropy(
-            pred_regs.view(-1, self.reg_max + 1),
-            gt_left.view(-1),
-            reduction='none').view(gt_left.shape) * weight_left
-        # loss right
-        loss_right = F.cross_entropy(
-            pred_regs.view(-1, self.reg_max + 1),
-            gt_right.view(-1),
-            reduction='none').view(gt_left.shape) * weight_right
-
-        loss = (loss_left + loss_right).mean(-1, keepdim=True)
-        
-        return loss
 
 
-    def forward(self, pred_regs, pred_boxs, anchors, gt_boxs, bbox_weight, fg_masks, strides):
+    def forward(self, pred_boxs, gt_boxs, bbox_weight, fg_masks):
         """
         Input:
-            pred_regs: (Tensor) [BM, 4*(reg_max + 1)]
             pred_boxs: (Tensor) [BM, 4]
             anchors: (Tensor) [BM, 2]
             gt_boxs: (Tensor) [BM, 4]
@@ -225,30 +258,18 @@ class RegressionLoss(nn.Module):
                             CIoU=True)
             loss_iou = (1.0 - ious) * bbox_weight
                
-            # dfl loss
-            if self.use_dfl:
-                pred_regs_pos = pred_regs[fg_masks]
-                gt_boxs_s = gt_boxs / strides
-                anchors_s = anchors / strides
-                gt_ltrb_s = bbox2dist(anchors_s, gt_boxs_s, self.reg_max)
-                gt_ltrb_s_pos = gt_ltrb_s[fg_masks]
-                loss_dfl = self.df_loss(pred_regs_pos, gt_ltrb_s_pos)
-                loss_dfl *= bbox_weight
-            else:
-                loss_dfl = pred_regs.sum() * 0.
-
         else:
-            loss_iou = pred_regs.sum() * 0.
-            loss_dfl = pred_regs.sum() * 0.
+            loss_iou = pred_boxs.sum() * 0.
 
-        return loss_iou, loss_dfl
+        return loss_iou
 
 
-def build_criterion(cfg, device, num_classes):
+def build_criterion(cfg, device, num_classes, warmup_epoch=1):
     criterion = Criterion(
         cfg=cfg,
         device=device,
-        num_classes=num_classes
+        num_classes=num_classes,
+        warmup_epoch=warmup_epoch,
         )
 
     return criterion

models/detectors/yolov7_plus/matcher.py

@@ -0,0 +1,419 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.box_ops import bbox_iou
+
+
+# -------------------------- YOLOv5 Assigner --------------------------
+class Yolov5Matcher(object):
+    def __init__(self, num_classes, num_anchors, anchor_size, anchor_theshold):
+        self.num_classes = num_classes
+        self.num_anchors = num_anchors
+        self.anchor_theshold = anchor_theshold
+        # [KA, 2]
+        self.anchor_sizes = np.array([[anchor[0], anchor[1]]
+                                      for anchor in anchor_size])
+        # [KA, 4]
+        self.anchor_boxes = np.array([[0., 0., anchor[0], anchor[1]]
+                                      for anchor in anchor_size])
+
+    def compute_iou(self, anchor_boxes, gt_box):
+        """
+            anchor_boxes : ndarray -> [KA, 4] (cx, cy, bw, bh).
+            gt_box : ndarray -> [1, 4] (cx, cy, bw, bh).
+        """
+        # anchors: [KA, 4]
+        anchors = np.zeros_like(anchor_boxes)
+        anchors[..., :2] = anchor_boxes[..., :2] - anchor_boxes[..., 2:] * 0.5  # x1y1
+        anchors[..., 2:] = anchor_boxes[..., :2] + anchor_boxes[..., 2:] * 0.5  # x2y2
+        anchors_area = anchor_boxes[..., 2] * anchor_boxes[..., 3]
+        
+        # gt_box: [1, 4] -> [KA, 4]
+        gt_box = np.array(gt_box).reshape(-1, 4)
+        gt_box = np.repeat(gt_box, anchors.shape[0], axis=0)
+        gt_box_ = np.zeros_like(gt_box)
+        gt_box_[..., :2] = gt_box[..., :2] - gt_box[..., 2:] * 0.5  # x1y1
+        gt_box_[..., 2:] = gt_box[..., :2] + gt_box[..., 2:] * 0.5  # x2y2
+        gt_box_area = np.prod(gt_box[..., 2:] - gt_box[..., :2], axis=1)
+
+        # intersection
+        inter_w = np.minimum(anchors[:, 2], gt_box_[:, 2]) - \
+                  np.maximum(anchors[:, 0], gt_box_[:, 0])
+        inter_h = np.minimum(anchors[:, 3], gt_box_[:, 3]) - \
+                  np.maximum(anchors[:, 1], gt_box_[:, 1])
+        inter_area = inter_w * inter_h
+        
+        # union
+        union_area = anchors_area + gt_box_area - inter_area
+
+        # iou
+        iou = inter_area / union_area
+        iou = np.clip(iou, a_min=1e-10, a_max=1.0)
+        
+        return iou
+
+
+    def iou_assignment(self, ctr_points, gt_box, fpn_strides):
+        # compute IoU
+        iou = self.compute_iou(self.anchor_boxes, gt_box)
+        iou_mask = (iou > 0.5)
+
+        label_assignment_results = []
+        if iou_mask.sum() == 0:
+            # We assign the anchor box with highest IoU score.
+            iou_ind = np.argmax(iou)
+
+            level = iou_ind // self.num_anchors              # pyramid level
+            anchor_idx = iou_ind - level * self.num_anchors  # anchor index
+
+            # get the corresponding stride
+            stride = fpn_strides[level]
+
+            # compute the grid cell
+            xc, yc = ctr_points
+            xc_s = xc / stride
+            yc_s = yc / stride
+            grid_x = int(xc_s)
+            grid_y = int(yc_s)
+
+            label_assignment_results.append([grid_x, grid_y, xc_s, yc_s, level, anchor_idx])
+        else:            
+            for iou_ind, iou_m in enumerate(iou_mask):
+                if iou_m:
+                    level = iou_ind // self.num_anchors              # pyramid level
+                    anchor_idx = iou_ind - level * self.num_anchors  # anchor index
+
+                    # get the corresponding stride
+                    stride = fpn_strides[level]
+
+                    # compute the gride cell
+                    xc, yc = ctr_points
+                    xc_s = xc / stride
+                    yc_s = yc / stride
+                    grid_x = int(xc_s)
+                    grid_y = int(yc_s)
+
+                    label_assignment_results.append([grid_x, grid_y, xc_s, yc_s, level, anchor_idx])
+
+        return label_assignment_results
+
+
+    def aspect_ratio_assignment(self, ctr_points, keeps, fpn_strides):
+        label_assignment_results = []
+        for keep_idx, keep in enumerate(keeps):
+            if keep:
+                level = keep_idx // self.num_anchors              # pyramid level
+                anchor_idx = keep_idx - level * self.num_anchors  # anchor index
+
+                # get the corresponding stride
+                stride = fpn_strides[level]
+
+                # compute the gride cell
+                xc, yc = ctr_points
+                xc_s = xc / stride
+                yc_s = yc / stride
+                grid_x = int(xc_s)
+                grid_y = int(yc_s)
+
+                label_assignment_results.append([grid_x, grid_y, xc_s, yc_s, level, anchor_idx])
+        
+        return label_assignment_results
+    
+
+    @torch.no_grad()
+    def __call__(self, fmp_sizes, fpn_strides, targets):
+        """
+            fmp_size: (List) [fmp_h, fmp_w]
+            fpn_strides: (List) -> [8, 16, 32, ...] stride of network output.
+            targets: (Dict) dict{'boxes': [...], 
+                                 'labels': [...], 
+                                 'orig_size': ...}
+        """
+        assert len(fmp_sizes) == len(fpn_strides)
+        # prepare
+        bs = len(targets)
+        gt_objectness = [
+            torch.zeros([bs, fmp_h, fmp_w, self.num_anchors, 1]) 
+            for (fmp_h, fmp_w) in fmp_sizes
+            ]
+        gt_classes = [
+            torch.zeros([bs, fmp_h, fmp_w, self.num_anchors, self.num_classes]) 
+            for (fmp_h, fmp_w) in fmp_sizes
+            ]
+        gt_bboxes = [
+            torch.zeros([bs, fmp_h, fmp_w, self.num_anchors, 4]) 
+            for (fmp_h, fmp_w) in fmp_sizes
+            ]
+
+        for batch_index in range(bs):
+            targets_per_image = targets[batch_index]
+            # [N,]
+            tgt_cls = targets_per_image["labels"].numpy()
+            # [N, 4]
+            tgt_box = targets_per_image['boxes'].numpy()
+
+            for gt_box, gt_label in zip(tgt_box, tgt_cls):
+                # get a bbox coords
+                x1, y1, x2, y2 = gt_box.tolist()
+                # xyxy -> cxcywh
+                xc, yc = (x2 + x1) * 0.5, (y2 + y1) * 0.5
+                bw, bh = x2 - x1, y2 - y1
+                gt_box = np.array([[0., 0., bw, bh]])
+
+                # check target
+                if bw < 1. or bh < 1.:
+                    # invalid target
+                    continue
+
+                # compute aspect ratio
+                ratios = gt_box[..., 2:] / self.anchor_sizes
+                keeps = np.maximum(ratios, 1 / ratios).max(-1) < self.anchor_theshold
+
+                if keeps.sum() == 0:
+                    label_assignment_results = self.iou_assignment([xc, yc], gt_box, fpn_strides)
+                else:
+                    label_assignment_results = self.aspect_ratio_assignment([xc, yc], keeps, fpn_strides)
+
+                # label assignment
+                for result in label_assignment_results:
+                    # assignment
+                    grid_x, grid_y, xc_s, yc_s, level, anchor_idx = result
+                    stride = fpn_strides[level]
+                    fmp_h, fmp_w = fmp_sizes[level]
+                    # coord on the feature
+                    x1s, y1s = x1 / stride, y1 / stride
+                    x2s, y2s = x2 / stride, y2 / stride
+                    # offset
+                    off_x = xc_s - grid_x
+                    off_y = yc_s - grid_y
+ 
+                    if off_x <= 0.5 and off_y <= 0.5:  # top left
+                        grids = [(grid_x-1, grid_y), (grid_x, grid_y-1), (grid_x, grid_y)]
+                    elif off_x > 0.5 and off_y <= 0.5: # top right
+                        grids = [(grid_x+1, grid_y), (grid_x, grid_y-1), (grid_x, grid_y)]
+                    elif off_x <= 0.5 and off_y > 0.5: # bottom left
+                        grids = [(grid_x-1, grid_y), (grid_x, grid_y+1), (grid_x, grid_y)]
+                    elif off_x > 0.5 and off_y > 0.5:  # bottom right
+                        grids = [(grid_x+1, grid_y), (grid_x, grid_y+1), (grid_x, grid_y)]
+
+                    for (i, j) in grids:
+                        is_in_box = (j >= y1s and j < y2s) and (i >= x1s and i < x2s)
+                        is_valid = (j >= 0 and j < fmp_h) and (i >= 0 and i < fmp_w)
+
+                        if is_in_box and is_valid:
+                            # obj
+                            gt_objectness[level][batch_index, j, i, anchor_idx] = 1.0
+                            # cls
+                            cls_ont_hot = torch.zeros(self.num_classes)
+                            cls_ont_hot[int(gt_label)] = 1.0
+                            gt_classes[level][batch_index, j, i, anchor_idx] = cls_ont_hot
+                            # box
+                            gt_bboxes[level][batch_index, j, i, anchor_idx] = torch.as_tensor([x1, y1, x2, y2])
+
+        # [B, M, C]
+        gt_objectness = torch.cat([gt.view(bs, -1, 1) for gt in gt_objectness], dim=1).float()
+        gt_classes = torch.cat([gt.view(bs, -1, self.num_classes) for gt in gt_classes], dim=1).float()
+        gt_bboxes = torch.cat([gt.view(bs, -1, 4) for gt in gt_bboxes], dim=1).float()
+
+        return gt_objectness, gt_classes, gt_bboxes
+
+
+# -------------------------- Task Aligned Assigner --------------------------
+class TaskAlignedAssigner(nn.Module):
+    def __init__(self,
+                 topk=10,
+                 num_classes=80,
+                 alpha=0.5,
+                 beta=6.0, 
+                 eps=1e-9):
+        super(TaskAlignedAssigner, self).__init__()
+        self.topk = topk
+        self.num_classes = num_classes
+        self.bg_idx = num_classes
+        self.alpha = alpha
+        self.beta = beta
+        self.eps = eps
+
+    @torch.no_grad()
+    def forward(self,
+                pd_scores,
+                pd_bboxes,
+                anc_points,
+                gt_labels,
+                gt_bboxes):
+        """This code referenced to
+           https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/tal_assigner.py
+        Args:
+            pd_scores (Tensor): shape(bs, num_total_anchors, num_classes)
+            pd_bboxes (Tensor): shape(bs, num_total_anchors, 4)
+            anc_points (Tensor): shape(num_total_anchors, 2)
+            gt_labels (Tensor): shape(bs, n_max_boxes, 1)
+            gt_bboxes (Tensor): shape(bs, n_max_boxes, 4)
+        Returns:
+            target_labels (Tensor): shape(bs, num_total_anchors)
+            target_bboxes (Tensor): shape(bs, num_total_anchors, 4)
+            target_scores (Tensor): shape(bs, num_total_anchors, num_classes)
+            fg_mask (Tensor): shape(bs, num_total_anchors)
+        """
+        self.bs = pd_scores.size(0)
+        self.n_max_boxes = gt_bboxes.size(1)
+
+        mask_pos, align_metric, overlaps = self.get_pos_mask(
+            pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points)
+
+        target_gt_idx, fg_mask, mask_pos = select_highest_overlaps(
+            mask_pos, overlaps, self.n_max_boxes)
+
+        # assigned target
+        target_labels, target_bboxes, target_scores = self.get_targets(
+            gt_labels, gt_bboxes, target_gt_idx, fg_mask)
+
+        # normalize
+        align_metric *= mask_pos
+        pos_align_metrics = align_metric.amax(axis=-1, keepdim=True)  # b, max_num_obj
+        pos_overlaps = (overlaps * mask_pos).amax(axis=-1, keepdim=True)  # b, max_num_obj
+        norm_align_metric = (align_metric * pos_overlaps / (pos_align_metrics + self.eps)).amax(-2).unsqueeze(-1)
+        target_scores = target_scores * norm_align_metric
+
+        return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
+
+
+    def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points):
+        # get anchor_align metric, (b, max_num_obj, h*w)
+        align_metric, overlaps = self.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes)
+        # get in_gts mask, (b, max_num_obj, h*w)
+        mask_in_gts = select_candidates_in_gts(anc_points, gt_bboxes)
+        # get topk_metric mask, (b, max_num_obj, h*w)
+        mask_topk = self.select_topk_candidates(align_metric * mask_in_gts)
+        # merge all mask to a final mask, (b, max_num_obj, h*w)
+        mask_pos = mask_topk * mask_in_gts
+
+        return mask_pos, align_metric, overlaps
+
+
+    def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes):
+        ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long)  # 2, b, max_num_obj
+        ind[0] = torch.arange(end=self.bs).view(-1, 1).repeat(1, self.n_max_boxes)  # b, max_num_obj
+        ind[1] = gt_labels.long().squeeze(-1)  # b, max_num_obj
+        # get the scores of each grid for each gt cls
+        bbox_scores = pd_scores[ind[0], :, ind[1]]  # b, max_num_obj, h*w
+
+        overlaps = bbox_iou(gt_bboxes.unsqueeze(2), pd_bboxes.unsqueeze(1), xywh=False,
+                            CIoU=True).squeeze(3).clamp(0)
+        align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta)
+
+        return align_metric, overlaps
+
+
+    def select_topk_candidates(self, metrics, largest=True):
+        """
+        Args:
+            metrics: (b, max_num_obj, h*w).
+            topk_mask: (b, max_num_obj, topk) or None
+        """
+
+        num_anchors = metrics.shape[-1]  # h*w
+        # (b, max_num_obj, topk)
+        topk_metrics, topk_idxs = torch.topk(metrics, self.topk, dim=-1, largest=largest)
+        topk_mask = (topk_metrics.max(-1, keepdim=True)[0] > self.eps).tile([1, 1, self.topk])
+        # (b, max_num_obj, topk)
+        topk_idxs[~topk_mask] = 0
+        # (b, max_num_obj, topk, h*w) -> (b, max_num_obj, h*w)
+        is_in_topk = F.one_hot(topk_idxs, num_anchors).sum(-2)
+        # filter invalid bboxes
+        is_in_topk = torch.where(is_in_topk > 1, 0, is_in_topk)
+        return is_in_topk.to(metrics.dtype)
+
+
+    def get_targets(self, gt_labels, gt_bboxes, target_gt_idx, fg_mask):
+        """
+        Args:
+            gt_labels: (b, max_num_obj, 1)
+            gt_bboxes: (b, max_num_obj, 4)
+            target_gt_idx: (b, h*w)
+            fg_mask: (b, h*w)
+        """
+
+        # assigned target labels, (b, 1)
+        batch_ind = torch.arange(end=self.bs, dtype=torch.int64, device=gt_labels.device)[..., None]
+        target_gt_idx = target_gt_idx + batch_ind * self.n_max_boxes  # (b, h*w)
+        target_labels = gt_labels.long().flatten()[target_gt_idx]  # (b, h*w)
+
+        # assigned target boxes, (b, max_num_obj, 4) -> (b, h*w)
+        target_bboxes = gt_bboxes.view(-1, 4)[target_gt_idx]
+
+        # assigned target scores
+        target_labels.clamp(0)
+        target_scores = F.one_hot(target_labels, self.num_classes)  # (b, h*w, 80)
+        fg_scores_mask = fg_mask[:, :, None].repeat(1, 1, self.num_classes)  # (b, h*w, 80)
+        target_scores = torch.where(fg_scores_mask > 0, target_scores, 0)
+
+        return target_labels, target_bboxes, target_scores
+    
+
+# -------------------------- Basic Functions --------------------------
+def select_candidates_in_gts(xy_centers, gt_bboxes, eps=1e-9):
+    """select the positive anchors's center in gt
+    Args:
+        xy_centers (Tensor): shape(bs*n_max_boxes, num_total_anchors, 4)
+        gt_bboxes (Tensor): shape(bs, n_max_boxes, 4)
+    Return:
+        (Tensor): shape(bs, n_max_boxes, num_total_anchors)
+    """
+    n_anchors = xy_centers.size(0)
+    bs, n_max_boxes, _ = gt_bboxes.size()
+    _gt_bboxes = gt_bboxes.reshape([-1, 4])
+    xy_centers = xy_centers.unsqueeze(0).repeat(bs * n_max_boxes, 1, 1)
+    gt_bboxes_lt = _gt_bboxes[:, 0:2].unsqueeze(1).repeat(1, n_anchors, 1)
+    gt_bboxes_rb = _gt_bboxes[:, 2:4].unsqueeze(1).repeat(1, n_anchors, 1)
+    b_lt = xy_centers - gt_bboxes_lt
+    b_rb = gt_bboxes_rb - xy_centers
+    bbox_deltas = torch.cat([b_lt, b_rb], dim=-1)
+    bbox_deltas = bbox_deltas.reshape([bs, n_max_boxes, n_anchors, -1])
+    return (bbox_deltas.min(axis=-1)[0] > eps).to(gt_bboxes.dtype)
+
+
+def select_highest_overlaps(mask_pos, overlaps, n_max_boxes):
+    """if an anchor box is assigned to multiple gts,
+        the one with the highest iou will be selected.
+    Args:
+        mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors)
+        overlaps (Tensor): shape(bs, n_max_boxes, num_total_anchors)
+    Return:
+        target_gt_idx (Tensor): shape(bs, num_total_anchors)
+        fg_mask (Tensor): shape(bs, num_total_anchors)
+        mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors)
+    """
+    fg_mask = mask_pos.sum(axis=-2)
+    if fg_mask.max() > 1:
+        mask_multi_gts = (fg_mask.unsqueeze(1) > 1).repeat([1, n_max_boxes, 1])
+        max_overlaps_idx = overlaps.argmax(axis=1)
+        is_max_overlaps = F.one_hot(max_overlaps_idx, n_max_boxes)
+        is_max_overlaps = is_max_overlaps.permute(0, 2, 1).to(overlaps.dtype)
+        mask_pos = torch.where(mask_multi_gts, is_max_overlaps, mask_pos)
+        fg_mask = mask_pos.sum(axis=-2)
+    target_gt_idx = mask_pos.argmax(axis=-2)
+    return target_gt_idx, fg_mask , mask_pos
+
+
+def iou_calculator(box1, box2, eps=1e-9):
+    """Calculate iou for batch
+    Args:
+        box1 (Tensor): shape(bs, n_max_boxes, 1, 4)
+        box2 (Tensor): shape(bs, 1, num_total_anchors, 4)
+    Return:
+        (Tensor): shape(bs, n_max_boxes, num_total_anchors)
+    """
+    box1 = box1.unsqueeze(2)  # [N, M1, 4] -> [N, M1, 1, 4]
+    box2 = box2.unsqueeze(1)  # [N, M2, 4] -> [N, 1, M2, 4]
+    px1y1, px2y2 = box1[:, :, :, 0:2], box1[:, :, :, 2:4]
+    gx1y1, gx2y2 = box2[:, :, :, 0:2], box2[:, :, :, 2:4]
+    x1y1 = torch.maximum(px1y1, gx1y1)
+    x2y2 = torch.minimum(px2y2, gx2y2)
+    overlap = (x2y2 - x1y1).clip(0).prod(-1)
+    area1 = (px2y2 - px1y1).clip(0).prod(-1)
+    area2 = (gx2y2 - gx1y1).clip(0).prod(-1)
+    union = area1 + area2 - overlap + eps
+
+    return overlap / union

models/detectors/yolov8/yolov8.py → models/detectors/yolov7_plus/yolov7_plus.py

@@ -2,16 +2,16 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .yolov8_backbone import build_backbone
-from .yolov8_neck import build_neck
-from .yolov8_pafpn import build_fpn
-from .yolov8_head import build_head
+from .yolov7_plus_backbone import build_backbone
+from .yolov7_plus_neck import build_neck
+from .yolov7_plus_pafpn import build_fpn
+from .yolov7_plus_head import build_head
 
 from utils.misc import multiclass_nms
 
 
 # Anchor-free YOLO
-class YOLOv8(nn.Module):
+class YOLOv7_Plus(nn.Module):
     def __init__(self, 
                  cfg,
                  device, 
@@ -20,22 +20,25 @@ class YOLOv8(nn.Module):
                  nms_thresh = 0.6,
                  trainable = False, 
                  topk = 1000):
-        super(YOLOv8, self).__init__()
+        super(YOLOv7_Plus, self).__init__()
         # --------- Basic Parameters ----------
         self.cfg = cfg
         self.device = device
         self.stride = cfg['stride']
-        self.reg_max = cfg['reg_max']
-        self.use_dfl = cfg['reg_max'] > 1
         self.num_classes = num_classes
         self.trainable = trainable
         self.conf_thresh = conf_thresh
         self.nms_thresh = nms_thresh
         self.topk = topk
         
-        # --------- Network Parameters ----------
-        self.proj_conv = nn.Conv2d(self.reg_max, 1, kernel_size=1, bias=False)
-
+        # ------------------- Anchor box -------------------
+        self.num_levels = 3
+        self.num_anchors = len(cfg['anchor_size']) // self.num_levels
+        self.anchor_size = torch.as_tensor(
+            cfg['anchor_size']
+            ).view(self.num_levels, self.num_anchors, 2) # [S, A, 2]
+        
+        # ------------------- Network Structure -------------------
         ## backbone
         self.backbone, feats_dim = build_backbone(cfg, cfg['pretrained']*trainable)
 
@@ -44,22 +47,22 @@ class YOLOv8(nn.Module):
         feats_dim[-1] = self.neck.out_dim
         
         ## fpn
-        self.fpn = build_fpn(cfg=cfg, in_dims=feats_dim)
-        fpn_dims = self.fpn.out_dim
+        self.fpn = build_fpn(cfg=cfg, in_dims=feats_dim, out_dim=round(256*cfg['width']))
+        self.head_dim = self.fpn.out_dim
 
         ## non-shared heads
         self.non_shared_heads = nn.ModuleList(
-            [build_head(cfg, feat_dim, fpn_dims, num_classes) 
-            for feat_dim in fpn_dims
+            [build_head(cfg, head_dim, head_dim, num_classes) 
+            for head_dim in self.head_dim
             ])
 
         ## pred
         self.cls_preds = nn.ModuleList(
-                            [nn.Conv2d(head.cls_out_dim, self.num_classes, kernel_size=1) 
+                            [nn.Conv2d(head.cls_out_dim, self.num_classes * self.num_anchors, kernel_size=1) 
                                 for head in self.non_shared_heads
                               ]) 
         self.reg_preds = nn.ModuleList(
-                            [nn.Conv2d(head.reg_out_dim, 4*(cfg['reg_max']), kernel_size=1) 
+                            [nn.Conv2d(head.reg_out_dim, 4 * self.num_anchors, kernel_size=1) 
                                 for head in self.non_shared_heads
                               ])                 
 
@@ -67,16 +70,24 @@ class YOLOv8(nn.Module):
     # ---------------------- Basic Functions ----------------------
     ## generate anchor points
     def generate_anchors(self, level, fmp_size):
-        """
-            fmp_size: (List) [H, W]
-        """
-        # generate grid cells
         fmp_h, fmp_w = fmp_size
+        # [KA, 2]
+        anchor_size = self.anchor_size[level]
+
+        # generate grid cells
         anchor_y, anchor_x = torch.meshgrid([torch.arange(fmp_h), torch.arange(fmp_w)])
-        # [H, W, 2] -> [HW, 2]
-        anchor_xy = torch.stack([anchor_x, anchor_y], dim=-1).float().view(-1, 2) + 0.5
+        anchor_xy = torch.stack([anchor_x, anchor_y], dim=-1).float().view(-1, 2)
+        # [HW, 2] -> [HW, KA, 2] -> [M, 2]
+        anchor_xy = anchor_xy.unsqueeze(1).repeat(1, self.num_anchors, 1)
+        anchor_xy = anchor_xy.view(-1, 2).to(self.device)
+        anchor_xy += 0.5
         anchor_xy *= self.stride[level]
-        anchors = anchor_xy.to(self.device)
+
+        # [KA, 2] -> [1, KA, 2] -> [HW, KA, 2] -> [M, 2]
+        anchor_wh = anchor_size.unsqueeze(0).repeat(fmp_h*fmp_w, 1, 1)
+        anchor_wh = anchor_wh.view(-1, 2).to(self.device)
+
+        anchors = torch.cat([anchor_xy, anchor_wh], dim=-1)
 
         return anchors
         
@@ -86,14 +97,13 @@ class YOLOv8(nn.Module):
         Input:
             cls_preds: List(Tensor) [[H x W, C], ...]
             box_preds: List(Tensor) [[H x W, 4], ...]
-            anchors:   List(Tensor) [[H x W, 2], ...]
         """
         all_scores = []
         all_labels = []
         all_bboxes = []
         
         for cls_pred_i, box_pred_i in zip(cls_preds, box_preds):
-            # (H x W x KA x C,)
+            # (H x W x C,)
             scores_i = cls_pred_i.sigmoid().flatten()
 
             # Keep top k top scoring indices only.
@@ -106,17 +116,16 @@ class YOLOv8(nn.Module):
 
             # filter out the proposals with low confidence score
             keep_idxs = topk_scores > self.conf_thresh
-            scores = topk_scores[keep_idxs]
+            topk_scores = topk_scores[keep_idxs]
             topk_idxs = topk_idxs[keep_idxs]
 
             anchor_idxs = torch.div(topk_idxs, self.num_classes, rounding_mode='floor')
-            labels = topk_idxs % self.num_classes
-
-            bboxes = box_pred_i[anchor_idxs]
+            topk_labels = topk_idxs % self.num_classes
+            topk_bboxes = box_pred_i[anchor_idxs]
 
-            all_scores.append(scores)
-            all_labels.append(labels)
-            all_bboxes.append(bboxes)
+            all_scores.append(topk_scores)
+            all_labels.append(topk_labels)
+            all_bboxes.append(topk_bboxes)
 
         scores = torch.cat(all_scores)
         labels = torch.cat(all_labels)
@@ -132,7 +141,7 @@ class YOLOv8(nn.Module):
             scores, labels, bboxes, self.nms_thresh, self.num_classes, False)
 
         return bboxes, scores, labels
-
+    
 
     # ---------------------- Main Process for Inference ----------------------
     @torch.no_grad()
@@ -154,8 +163,8 @@ class YOLOv8(nn.Module):
             cls_feat, reg_feat = head(feat)
 
             # pred
-            cls_pred = self.cls_preds[level](cls_feat)  # [B, C, H, W]
-            reg_pred = self.reg_preds[level](reg_feat)  # [B, 4*(reg_max), H, W]
+            cls_pred = self.cls_preds[level](cls_feat)  # [B, C*A, H, W]
+            reg_pred = self.reg_preds[level](reg_feat)  # [B, 4*A, H, W]
 
             B, _, H, W = cls_pred.size()
             fmp_size = [H, W]
@@ -163,22 +172,14 @@ class YOLOv8(nn.Module):
             anchors = self.generate_anchors(level, fmp_size)
 
             # [B, C, H, W] -> [B, H, W, C] -> [B, M, C]
-            cls_pred = cls_pred.permute(0, 2, 3, 1).contiguous().view(B, -1, self.num_classes)
-            reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(B, -1, 4*self.reg_max)
+            cls_pred = cls_pred[0].permute(1, 2, 0).contiguous().view(-1, self.num_classes)
+            reg_pred = reg_pred[0].permute(1, 2, 0).contiguous().view(-1, 4)
 
             # decode bbox
-            if self.use_dfl:
-                B, M = reg_pred.shape[:2]
-                # [B, M, 4*(reg_max)] -> [B, M, 4, reg_max] -> [B, 4, M, reg_max]
-                reg_pred = reg_pred.reshape([B, M, 4, self.reg_max])
-                # [B, M, 4, reg_max] -> [B, reg_max, 4, M]
-                reg_pred = reg_pred.permute(0, 3, 2, 1).contiguous()
-                # [B, reg_max, 4, M] -> [B, 1, 4, M]
-                reg_pred = self.proj_conv(F.softmax(reg_pred, dim=1))
-                # [B, 1, 4, M] -> [B, 4, M] -> [B, M, 4]
-                reg_pred = reg_pred.view(B, 4, M).permute(0, 2, 1).contiguous()
-            pred_x1y1 = anchors - reg_pred[..., :2] * self.stride[level]
-            pred_x2y2 = anchors + reg_pred[..., 2:] * self.stride[level]
+            ctr_pred = anchors[..., :2] + reg_pred[..., :2] * self.stride[level]
+            wh_pred = torch.exp(reg_pred[..., 2:]) * anchors[..., 2:]
+            pred_x1y1 = ctr_pred - wh_pred * 0.5
+            pred_x2y2 = ctr_pred + wh_pred * 0.5
             box_pred = torch.cat([pred_x1y1, pred_x2y2], dim=-1)
 
             all_cls_preds.append(cls_pred)
@@ -187,7 +188,7 @@ class YOLOv8(nn.Module):
 
         # post process
         bboxes, scores, labels = self.post_process(
-            all_cls_preds, all_box_preds, all_anchors)
+            all_cls_preds, all_box_preds)
         
         return bboxes, scores, labels
 
@@ -207,11 +208,10 @@ class YOLOv8(nn.Module):
             pyramid_feats = self.fpn(pyramid_feats)
 
             # non-shared heads
-            all_anchors = []
+            all_fmp_sizes = []
             all_cls_preds = []
-            all_reg_preds = []
             all_box_preds = []
-            all_strides = []
+            all_anchors = []
             for level, (feat, head) in enumerate(zip(pyramid_feats, self.non_shared_heads)):
                 cls_feat, reg_feat = head(feat)
 
@@ -226,40 +226,26 @@ class YOLOv8(nn.Module):
                 
                 # [B, C, H, W] -> [B, H, W, C] -> [B, M, C]
                 cls_pred = cls_pred.permute(0, 2, 3, 1).contiguous().view(B, -1, self.num_classes)
-                reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(B, -1, 4*self.reg_max)
+                reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(B, -1, 4)
 
                 # decode bbox
-                if self.use_dfl:
-                    B, M = reg_pred.shape[:2]
-                    # [B, M, 4*(reg_max)] -> [B, M, 4, reg_max] -> [B, 4, M, reg_max]
-                    reg_pred_ = reg_pred.reshape([B, M, 4, self.reg_max]).clone()
-                    # [B, M, 4, reg_max] -> [B, reg_max, 4, M]
-                    reg_pred_ = reg_pred_.permute(0, 3, 2, 1).contiguous()
-                    # [B, reg_max, 4, M] -> [B, 1, 4, M]
-                    reg_pred_ = self.proj_conv(F.softmax(reg_pred_, dim=1))
-                    # [B, 1, 4, M] -> [B, 4, M] -> [B, M, 4]
-                    reg_pred_ = reg_pred_.view(B, 4, M).permute(0, 2, 1).contiguous()
-                pred_x1y1 = anchors - reg_pred_[..., :2] * self.stride[level]
-                pred_x2y2 = anchors + reg_pred_[..., 2:] * self.stride[level]
+                ctr_pred = anchors[..., :2] + reg_pred[..., :2] * self.stride[level]
+                wh_pred = torch.exp(reg_pred[..., 2:]) * anchors[..., 2:]
+                pred_x1y1 = ctr_pred - wh_pred * 0.5
+                pred_x2y2 = ctr_pred + wh_pred * 0.5
                 box_pred = torch.cat([pred_x1y1, pred_x2y2], dim=-1)
 
-                del reg_pred_
-
-                # stride tensor: [M, 1]
-                stride_tensor = torch.ones_like(anchors[..., :1]) * self.stride[level]
-
+                all_fmp_sizes.append(fmp_size)
                 all_cls_preds.append(cls_pred)
-                all_reg_preds.append(reg_pred)
                 all_box_preds.append(box_pred)
                 all_anchors.append(anchors)
-                all_strides.append(stride_tensor)
             
             # output dict
             outputs = {"pred_cls": all_cls_preds,        # List(Tensor) [B, M, C]
-                       "pred_reg": all_reg_preds,        # List(Tensor) [B, M, 4*(reg_max)]
                        "pred_box": all_box_preds,        # List(Tensor) [B, M, 4]
-                       "anchors": all_anchors,           # List(Tensor) [M, 2]
-                       "strides": self.stride,           # List(Int) = [8, 16, 32]
-                       "stride_tensor": all_strides      # List(Tensor) [M, 1]
-                       }           
-            return outputs
+                       "anchors": all_anchors,           # List(Tensor) [B, M, 4]
+                       'fmp_sizes': all_fmp_sizes,       # List
+                       'strides': self.stride,           # List
+                       }
+
+            return outputs 

models/detectors/yolov8/yolov8_backbone.py → models/detectors/yolov7_plus/yolov7_plus_backbone.py

@@ -2,16 +2,16 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov8_basic import Conv, ELAN_CSP_Block
+    from .yolov7_plus_basic import Conv, ELAN_CSP_Block
 except:
-    from yolov8_basic import Conv, ELAN_CSP_Block
+    from yolov7_plus_basic import Conv, ELAN_CSP_Block
 
 
 # ---------------------------- ImageNet pretrained weights ----------------------------
 model_urls = {
     'elan_cspnet_nano': "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_nano.pth",
-    'elan_cspnet_small': None,
-    'elan_cspnet_medium': None,
+    'elan_cspnet_small': "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_small.pth",
+    'elan_cspnet_medium': "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_medium.pth",
     'elan_cspnet_large': "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_large.pth",
     'elan_cspnet_huge': None,
 }

models/detectors/yolov8/yolov8_basic.py → models/detectors/yolov7_plus/yolov7_plus_basic.py

@@ -130,3 +130,34 @@ class ELAN_CSP_Block(nn.Module):
         out = self.cv3(torch.cat(out, dim=1))
 
         return out
+
+
+# ---------------------------- FPN Modules ----------------------------
+## build fpn's core block
+def build_fpn_block(cfg, in_dim, out_dim):
+    if cfg['fpn_core_block'] == 'ELAN_CSPBlock':
+        layer = ELAN_CSP_Block(in_dim=in_dim,
+                               out_dim=out_dim,
+                               expand_ratio=cfg['expand_ratio'],
+                               nblocks=round(3*cfg['depth']),
+                               shortcut=False,
+                               act_type=cfg['fpn_act'],
+                               norm_type=cfg['fpn_norm'],
+                               depthwise=cfg['fpn_depthwise']
+                               )
+        
+    return layer
+
+## build fpn's reduce layer
+def build_reduce_layer(cfg, in_dim, out_dim):
+    if cfg['fpn_reduce_layer'] == 'Conv':
+        layer = Conv(in_dim, out_dim, k=1, act_type=cfg['fpn_act'], norm_type=cfg['fpn_norm'])
+        
+    return layer
+
+## build fpn's downsample layer
+def build_downsample_layer(cfg, in_dim, out_dim):
+    if cfg['fpn_downsample_layer'] == 'Conv':
+        layer = Conv(in_dim, out_dim, k=3, s=2, p=1, act_type=cfg['fpn_act'], norm_type=cfg['fpn_norm'])
+        
+    return layer

models/detectors/yolov8/yolov8_head.py → models/detectors/yolov7_plus/yolov7_plus_head.py

@@ -1,13 +1,13 @@
 import torch
 import torch.nn as nn
 try:
-    from .yolov8_basic import Conv
+    from .yolov7_plus_basic import Conv
 except:
-    from yolov8_basic import Conv
+    from yolov7_plus_basic import Conv
 
 
 class DecoupledHead(nn.Module):
-    def __init__(self, cfg, in_dim, fpn_dims, num_classes=80):
+    def __init__(self, cfg, in_dim, out_dim, num_classes=80):
         super().__init__()
         print('==============================')
         print('Head: Decoupled Head')
@@ -19,7 +19,7 @@ class DecoupledHead(nn.Module):
 
         # cls head
         cls_feats = []
-        self.cls_out_dim = max(fpn_dims[0], num_classes)
+        self.cls_out_dim = max(out_dim, num_classes)
         for i in range(cfg['num_cls_head']):
             if i == 0:
                 cls_feats.append(
@@ -38,7 +38,7 @@ class DecoupledHead(nn.Module):
                 
         # reg head
         reg_feats = []
-        self.reg_out_dim = max(16, fpn_dims[0]//4, 4*cfg['reg_max'])
+        self.reg_out_dim = out_dim
         for i in range(cfg['num_reg_head']):
             if i == 0:
                 reg_feats.append(

models/detectors/yolov7_plus/yolov7_plus_neck.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+from .yolov7_plus_basic import Conv
+
+
+# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
+class SPPF(nn.Module):
+    """
+        This code referenced to https://github.com/ultralytics/yolov5
+    """
+    def __init__(self, cfg, in_dim, out_dim, expand_ratio=0.5):
+        super().__init__()
+        inter_dim = int(in_dim * expand_ratio)
+        self.out_dim = out_dim
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=cfg['neck_act'], norm_type=cfg['neck_norm'])
+        self.cv2 = Conv(inter_dim * 4, out_dim, k=1, act_type=cfg['neck_act'], norm_type=cfg['neck_norm'])
+        self.m = nn.MaxPool2d(kernel_size=cfg['pooling_size'], stride=1, padding=cfg['pooling_size'] // 2)
+
+    def forward(self, x):
+        x = self.cv1(x)
+        y1 = self.m(x)
+        y2 = self.m(y1)
+
+        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+# SPPF block with CSP module
+class SPPFBlockCSP(nn.Module):
+    """
+        CSP Spatial Pyramid Pooling Block
+    """
+    def __init__(self, cfg, in_dim, out_dim):
+        super(SPPFBlockCSP, self).__init__()
+        inter_dim = int(in_dim * cfg['expand_ratio'])
+        self.out_dim = out_dim
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=cfg['neck_act'], norm_type=cfg['neck_norm'])
+        self.cv2 = Conv(in_dim, inter_dim, k=1, act_type=cfg['neck_act'], norm_type=cfg['neck_norm'])
+        self.m = nn.Sequential(
+            Conv(inter_dim, inter_dim, k=3, p=1, 
+                 act_type=cfg['neck_act'], norm_type=cfg['neck_norm'], 
+                 depthwise=cfg['neck_depthwise']),
+            SPPF(cfg, inter_dim, inter_dim, expand_ratio=1.0),
+            Conv(inter_dim, inter_dim, k=3, p=1, 
+                 act_type=cfg['neck_act'], norm_type=cfg['neck_norm'], 
+                 depthwise=cfg['neck_depthwise'])
+        )
+        self.cv3 = Conv(inter_dim * 2, self.out_dim, k=1, act_type=cfg['neck_act'], norm_type=cfg['neck_norm'])
+
+        
+    def forward(self, x):
+        x1 = self.cv1(x)
+        x2 = self.cv2(x)
+        x3 = self.m(x2)
+        y = self.cv3(torch.cat([x1, x3], dim=1))
+
+        return y
+
+
+def build_neck(cfg, in_dim, out_dim):
+    model = cfg['neck']
+    print('==============================')
+    print('Neck: {}'.format(model))
+    # build neck
+    if model == 'sppf':
+        neck = SPPF(cfg, in_dim, out_dim, cfg['expand_ratio'])
+    elif model == 'csp_sppf':
+        neck = SPPFBlockCSP(cfg, in_dim, out_dim)
+
+    return neck
+        

models/detectors/yolov7_plus/yolov7_plus_pafpn.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .yolov7_plus_basic import (Conv, build_reduce_layer, build_downsample_layer, build_fpn_block)
+
+
+# YOLO-Style PaFPN
+class Yolov7PlusPaFPN(nn.Module):
+    def __init__(self, cfg, in_dims=[256, 512, 1024], out_dim=None):
+        super(Yolov7PlusPaFPN, self).__init__()
+        # --------------------------- Basic Parameters ---------------------------
+        self.in_dims = in_dims
+        c3, c4, c5 = in_dims
+        width = cfg['width']
+
+        # --------------------------- Network Parameters ---------------------------
+        ## top dwon
+        ### P5 -> P4
+        self.reduce_layer_1 = build_reduce_layer(cfg, c5, round(512*width))
+        self.top_down_layer_1 = build_fpn_block(cfg, c4 + round(512*width), round(512*width))
+
+        ### P4 -> P3
+        self.reduce_layer_2 = build_reduce_layer(cfg, round(512*width), round(256*width))
+        self.top_down_layer_2 = build_fpn_block(cfg, c3 + round(256*width), round(256*width))
+
+        ## bottom up
+        ### P3 -> P4
+        self.downsample_layer_1 = build_downsample_layer(cfg, round(256*width), round(256*width))
+        self.bottom_up_layer_1 = build_fpn_block(cfg, round(256*width) + round(256*width), round(512*width))
+
+        ### P4 -> P5
+        self.downsample_layer_2 = build_downsample_layer(cfg, round(512*width), round(512*width))
+        self.bottom_up_layer_2 = build_fpn_block(cfg, round(512*width) + round(512*width), round(1024*width))
+                
+        ## output proj layers
+        if out_dim is not None:
+            self.out_layers = nn.ModuleList([
+                Conv(in_dim, out_dim, k=1,
+                     act_type=cfg['fpn_act'], norm_type=cfg['fpn_norm'])
+                     for in_dim in [round(256*width), round(512*width), round(1024*width)]
+                     ])
+            self.out_dim = [out_dim] * 3
+        else:
+            self.out_layers = None
+            self.out_dim = [round(256*width), round(512*width), round(1024*width)]
+
+
+    def forward(self, features):
+        c3, c4, c5 = features
+
+        # Top down
+        ## P5 -> P4
+        c6 = self.reduce_layer_1(c5)
+        c7 = F.interpolate(c6, scale_factor=2.0)
+        c8 = torch.cat([c7, c4], dim=1)
+        c9 = self.top_down_layer_1(c8)
+        ## P4 -> P3
+        c10 = self.reduce_layer_2(c9)
+        c11 = F.interpolate(c10, scale_factor=2.0)
+        c12 = torch.cat([c11, c3], dim=1)
+        c13 = self.top_down_layer_2(c12)
+
+        # Bottom up
+        ## p3 -> P4
+        c14 = self.downsample_layer_1(c13)
+        c15 = torch.cat([c14, c10], dim=1)
+        c16 = self.bottom_up_layer_1(c15)
+        ## P4 -> P5
+        c17 = self.downsample_layer_2(c16)
+        c18 = torch.cat([c17, c6], dim=1)
+        c19 = self.bottom_up_layer_2(c18)
+
+        out_feats = [c13, c16, c19] # [P3, P4, P5]
+        
+        # output proj layers
+        if self.out_layers is not None:
+            out_feats_proj = []
+            for feat, layer in zip(out_feats, self.out_layers):
+                out_feats_proj.append(layer(feat))
+            return out_feats_proj
+
+        return out_feats
+
+
+def build_fpn(cfg, in_dims, out_dim=None):
+    model = cfg['fpn']
+    # build pafpn
+    if model == 'yolov7_plus_pafpn':
+        fpn_net = Yolov7PlusPaFPN(cfg, in_dims, out_dim)
+
+    return fpn_net

models/detectors/yolov8/matcher.py

@@ -1,204 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from utils.box_ops import bbox_iou
-
-
-# -------------------------- Task Aligned Assigner --------------------------
-class TaskAlignedAssigner(nn.Module):
-    def __init__(self,
-                 topk=10,
-                 num_classes=80,
-                 alpha=0.5,
-                 beta=6.0, 
-                 eps=1e-9):
-        super(TaskAlignedAssigner, self).__init__()
-        self.topk = topk
-        self.num_classes = num_classes
-        self.bg_idx = num_classes
-        self.alpha = alpha
-        self.beta = beta
-        self.eps = eps
-
-    @torch.no_grad()
-    def forward(self,
-                pd_scores,
-                pd_bboxes,
-                anc_points,
-                gt_labels,
-                gt_bboxes):
-        """This code referenced to
-           https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/tal_assigner.py
-        Args:
-            pd_scores (Tensor): shape(bs, num_total_anchors, num_classes)
-            pd_bboxes (Tensor): shape(bs, num_total_anchors, 4)
-            anc_points (Tensor): shape(num_total_anchors, 2)
-            gt_labels (Tensor): shape(bs, n_max_boxes, 1)
-            gt_bboxes (Tensor): shape(bs, n_max_boxes, 4)
-        Returns:
-            target_labels (Tensor): shape(bs, num_total_anchors)
-            target_bboxes (Tensor): shape(bs, num_total_anchors, 4)
-            target_scores (Tensor): shape(bs, num_total_anchors, num_classes)
-            fg_mask (Tensor): shape(bs, num_total_anchors)
-        """
-        self.bs = pd_scores.size(0)
-        self.n_max_boxes = gt_bboxes.size(1)
-
-        mask_pos, align_metric, overlaps = self.get_pos_mask(
-            pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points)
-
-        target_gt_idx, fg_mask, mask_pos = select_highest_overlaps(
-            mask_pos, overlaps, self.n_max_boxes)
-
-        # assigned target
-        target_labels, target_bboxes, target_scores = self.get_targets(
-            gt_labels, gt_bboxes, target_gt_idx, fg_mask)
-
-        # normalize
-        align_metric *= mask_pos
-        pos_align_metrics = align_metric.amax(axis=-1, keepdim=True)  # b, max_num_obj
-        pos_overlaps = (overlaps * mask_pos).amax(axis=-1, keepdim=True)  # b, max_num_obj
-        norm_align_metric = (align_metric * pos_overlaps / (pos_align_metrics + self.eps)).amax(-2).unsqueeze(-1)
-        target_scores = target_scores * norm_align_metric
-
-        return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
-
-
-    def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points):
-        # get anchor_align metric, (b, max_num_obj, h*w)
-        align_metric, overlaps = self.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes)
-        # get in_gts mask, (b, max_num_obj, h*w)
-        mask_in_gts = select_candidates_in_gts(anc_points, gt_bboxes)
-        # get topk_metric mask, (b, max_num_obj, h*w)
-        mask_topk = self.select_topk_candidates(align_metric * mask_in_gts)
-        # merge all mask to a final mask, (b, max_num_obj, h*w)
-        mask_pos = mask_topk * mask_in_gts
-
-        return mask_pos, align_metric, overlaps
-
-
-    def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes):
-        ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long)  # 2, b, max_num_obj
-        ind[0] = torch.arange(end=self.bs).view(-1, 1).repeat(1, self.n_max_boxes)  # b, max_num_obj
-        ind[1] = gt_labels.long().squeeze(-1)  # b, max_num_obj
-        # get the scores of each grid for each gt cls
-        bbox_scores = pd_scores[ind[0], :, ind[1]]  # b, max_num_obj, h*w
-
-        overlaps = bbox_iou(gt_bboxes.unsqueeze(2), pd_bboxes.unsqueeze(1), xywh=False,
-                            CIoU=True).squeeze(3).clamp(0)
-        align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta)
-
-        return align_metric, overlaps
-
-
-    def select_topk_candidates(self, metrics, largest=True):
-        """
-        Args:
-            metrics: (b, max_num_obj, h*w).
-            topk_mask: (b, max_num_obj, topk) or None
-        """
-
-        num_anchors = metrics.shape[-1]  # h*w
-        # (b, max_num_obj, topk)
-        topk_metrics, topk_idxs = torch.topk(metrics, self.topk, dim=-1, largest=largest)
-        topk_mask = (topk_metrics.max(-1, keepdim=True)[0] > self.eps).tile([1, 1, self.topk])
-        # (b, max_num_obj, topk)
-        topk_idxs[~topk_mask] = 0
-        # (b, max_num_obj, topk, h*w) -> (b, max_num_obj, h*w)
-        is_in_topk = F.one_hot(topk_idxs, num_anchors).sum(-2)
-        # filter invalid bboxes
-        is_in_topk = torch.where(is_in_topk > 1, 0, is_in_topk)
-        return is_in_topk.to(metrics.dtype)
-
-
-    def get_targets(self, gt_labels, gt_bboxes, target_gt_idx, fg_mask):
-        """
-        Args:
-            gt_labels: (b, max_num_obj, 1)
-            gt_bboxes: (b, max_num_obj, 4)
-            target_gt_idx: (b, h*w)
-            fg_mask: (b, h*w)
-        """
-
-        # assigned target labels, (b, 1)
-        batch_ind = torch.arange(end=self.bs, dtype=torch.int64, device=gt_labels.device)[..., None]
-        target_gt_idx = target_gt_idx + batch_ind * self.n_max_boxes  # (b, h*w)
-        target_labels = gt_labels.long().flatten()[target_gt_idx]  # (b, h*w)
-
-        # assigned target boxes, (b, max_num_obj, 4) -> (b, h*w)
-        target_bboxes = gt_bboxes.view(-1, 4)[target_gt_idx]
-
-        # assigned target scores
-        target_labels.clamp(0)
-        target_scores = F.one_hot(target_labels, self.num_classes)  # (b, h*w, 80)
-        fg_scores_mask = fg_mask[:, :, None].repeat(1, 1, self.num_classes)  # (b, h*w, 80)
-        target_scores = torch.where(fg_scores_mask > 0, target_scores, 0)
-
-        return target_labels, target_bboxes, target_scores
-    
-
-# -------------------------- Basic Functions --------------------------
-def select_candidates_in_gts(xy_centers, gt_bboxes, eps=1e-9):
-    """select the positive anchors's center in gt
-    Args:
-        xy_centers (Tensor): shape(bs*n_max_boxes, num_total_anchors, 4)
-        gt_bboxes (Tensor): shape(bs, n_max_boxes, 4)
-    Return:
-        (Tensor): shape(bs, n_max_boxes, num_total_anchors)
-    """
-    n_anchors = xy_centers.size(0)
-    bs, n_max_boxes, _ = gt_bboxes.size()
-    _gt_bboxes = gt_bboxes.reshape([-1, 4])
-    xy_centers = xy_centers.unsqueeze(0).repeat(bs * n_max_boxes, 1, 1)
-    gt_bboxes_lt = _gt_bboxes[:, 0:2].unsqueeze(1).repeat(1, n_anchors, 1)
-    gt_bboxes_rb = _gt_bboxes[:, 2:4].unsqueeze(1).repeat(1, n_anchors, 1)
-    b_lt = xy_centers - gt_bboxes_lt
-    b_rb = gt_bboxes_rb - xy_centers
-    bbox_deltas = torch.cat([b_lt, b_rb], dim=-1)
-    bbox_deltas = bbox_deltas.reshape([bs, n_max_boxes, n_anchors, -1])
-    return (bbox_deltas.min(axis=-1)[0] > eps).to(gt_bboxes.dtype)
-
-
-def select_highest_overlaps(mask_pos, overlaps, n_max_boxes):
-    """if an anchor box is assigned to multiple gts,
-        the one with the highest iou will be selected.
-    Args:
-        mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors)
-        overlaps (Tensor): shape(bs, n_max_boxes, num_total_anchors)
-    Return:
-        target_gt_idx (Tensor): shape(bs, num_total_anchors)
-        fg_mask (Tensor): shape(bs, num_total_anchors)
-        mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors)
-    """
-    fg_mask = mask_pos.sum(axis=-2)
-    if fg_mask.max() > 1:
-        mask_multi_gts = (fg_mask.unsqueeze(1) > 1).repeat([1, n_max_boxes, 1])
-        max_overlaps_idx = overlaps.argmax(axis=1)
-        is_max_overlaps = F.one_hot(max_overlaps_idx, n_max_boxes)
-        is_max_overlaps = is_max_overlaps.permute(0, 2, 1).to(overlaps.dtype)
-        mask_pos = torch.where(mask_multi_gts, is_max_overlaps, mask_pos)
-        fg_mask = mask_pos.sum(axis=-2)
-    target_gt_idx = mask_pos.argmax(axis=-2)
-    return target_gt_idx, fg_mask , mask_pos
-
-
-def iou_calculator(box1, box2, eps=1e-9):
-    """Calculate iou for batch
-    Args:
-        box1 (Tensor): shape(bs, n_max_boxes, 1, 4)
-        box2 (Tensor): shape(bs, 1, num_total_anchors, 4)
-    Return:
-        (Tensor): shape(bs, n_max_boxes, num_total_anchors)
-    """
-    box1 = box1.unsqueeze(2)  # [N, M1, 4] -> [N, M1, 1, 4]
-    box2 = box2.unsqueeze(1)  # [N, M2, 4] -> [N, 1, M2, 4]
-    px1y1, px2y2 = box1[:, :, :, 0:2], box1[:, :, :, 2:4]
-    gx1y1, gx2y2 = box2[:, :, :, 0:2], box2[:, :, :, 2:4]
-    x1y1 = torch.maximum(px1y1, gx1y1)
-    x2y2 = torch.minimum(px2y2, gx2y2)
-    overlap = (x2y2 - x1y1).clip(0).prod(-1)
-    area1 = (px2y2 - px1y1).clip(0).prod(-1)
-    area2 = (gx2y2 - gx1y1).clip(0).prod(-1)
-    union = area1 + area2 - overlap + eps
-
-    return overlap / union

models/detectors/yolov8/yolov8_neck.py

@@ -1,101 +0,0 @@
-import torch
-import torch.nn as nn
-
-try:
-    from .yolov8_basic import Conv
-except:
-    from yolov8_basic import Conv
-
-# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
-class SPPF(nn.Module):
-    """
-        This code referenced to https://github.com/ultralytics/yolov5
-    """
-    def __init__(self, in_dim, out_dim, expand_ratio=0.5, pooling_size=5, act_type='', norm_type=''):
-        super().__init__()
-        inter_dim = int(in_dim * expand_ratio)
-        self.out_dim = out_dim
-        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
-        self.cv2 = Conv(inter_dim * 4, out_dim, k=1, act_type=act_type, norm_type=norm_type)
-        self.m = nn.MaxPool2d(kernel_size=pooling_size, stride=1, padding=pooling_size // 2)
-
-    def forward(self, x):
-        x = self.cv1(x)
-        y1 = self.m(x)
-        y2 = self.m(y1)
-
-        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
-
-
-# SPPF block with CSP module
-class SPPFBlockCSP(nn.Module):
-    """
-        CSP Spatial Pyramid Pooling Block
-    """
-    def __init__(self,
-                 in_dim,
-                 out_dim,
-                 expand_ratio=0.5,
-                 pooling_size=5,
-                 act_type='lrelu',
-                 norm_type='BN',
-                 depthwise=False
-                 ):
-        super(SPPFBlockCSP, self).__init__()
-        inter_dim = int(in_dim * expand_ratio)
-        self.out_dim = out_dim
-        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
-        self.cv2 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
-        self.m = nn.Sequential(
-            Conv(inter_dim, inter_dim, k=3, p=1, 
-                 act_type=act_type, norm_type=norm_type, 
-                 depthwise=depthwise),
-            SPPF(inter_dim, 
-                 inter_dim, 
-                 expand_ratio=1.0, 
-                 pooling_size=pooling_size, 
-                 act_type=act_type, 
-                 norm_type=norm_type),
-            Conv(inter_dim, inter_dim, k=3, p=1, 
-                 act_type=act_type, norm_type=norm_type, 
-                 depthwise=depthwise)
-        )
-        self.cv3 = Conv(inter_dim * 2, self.out_dim, k=1, act_type=act_type, norm_type=norm_type)
-
-        
-    def forward(self, x):
-        x1 = self.cv1(x)
-        x2 = self.cv2(x)
-        x3 = self.m(x2)
-        y = self.cv3(torch.cat([x1, x3], dim=1))
-
-        return y
-
-
-def build_neck(cfg, in_dim, out_dim):
-    model = cfg['neck']
-    print('==============================')
-    print('Neck: {}'.format(model))
-    # build neck
-    if model == 'sppf':
-        neck = SPPF(
-            in_dim=in_dim,
-            out_dim=out_dim,
-            expand_ratio=cfg['expand_ratio'], 
-            pooling_size=cfg['pooling_size'],
-            act_type=cfg['neck_act'],
-            norm_type=cfg['neck_norm']
-            )
-    elif model == 'sppf_block_csp':
-        neck = SPPFBlockCSP(
-            in_dim=in_dim,
-            out_dim=out_dim,
-            expand_ratio=cfg['expand_ratio'], 
-            pooling_size=cfg['pooling_size'],
-            act_type=cfg['neck_act'],
-            norm_type=cfg['neck_norm'],
-            depthwise=cfg['neck_depthwise']
-            )
-
-    return neck
-    

models/detectors/yolov8/yolov8_pafpn.py

@@ -1,145 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-try:
-    from .yolov8_basic import Conv, ELAN_CSP_Block
-except:
-    from yolov8_basic import Conv, ELAN_CSP_Block
-
-
-# PaFPN-ELAN
-class Yolov8PaFPN(nn.Module):
-    def __init__(self, 
-                 in_dims=[256, 512, 512],
-                 width=1.0,
-                 depth=1.0,
-                 ratio=1.0,
-                 act_type='silu',
-                 norm_type='BN',
-                 depthwise=False):
-        super(Yolov8PaFPN, self).__init__()
-        print('==============================')
-        print('FPN: {}'.format("ELAN_PaFPN"))
-        self.in_dims = in_dims
-        self.width = width
-        self.depth = depth
-        self.out_dim = [int(256 * width), int(512 * width), int(512 * width * ratio)]
-        c3, c4, c5 = in_dims
-
-        # --------------------------- Top-dwon ---------------------------
-        ## P5 -> P4
-        self.head_elan_1 = ELAN_CSP_Block(in_dim=c5 + c4,
-                                          out_dim=int(512*width),
-                                          expand_ratio=0.5,
-                                          nblocks=int(3*depth),
-                                          shortcut=False,
-                                          depthwise=depthwise,
-                                          norm_type=norm_type,
-                                          act_type=act_type
-                                          )
-        ## P4 -> P3
-        self.head_elan_2 = ELAN_CSP_Block(in_dim=int(512*width) + c3,
-                                          out_dim=int(256*width),
-                                          expand_ratio=0.5,
-                                          nblocks=int(3*depth),
-                                          shortcut=False,
-                                          depthwise=depthwise,
-                                          norm_type=norm_type,
-                                          act_type=act_type
-                                          )
-        # --------------------------- Bottom-up ---------------------------
-        ## P3 -> P4
-        self.mp1 = Conv(int(256*width), int(256*width), k=3, p=1, s=2,
-                        act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.head_elan_3 = ELAN_CSP_Block(in_dim=int(256*width) + int(512*width),
-                                          out_dim=int(512*width),
-                                          expand_ratio=0.5,
-                                          nblocks=int(3*depth),
-                                          shortcut=False,
-                                          depthwise=depthwise,
-                                          norm_type=norm_type,
-                                          act_type=act_type
-                                          )
-        ## P4 -> P5
-        self.mp2 = Conv(int(512 * width), int(512 * width), k=3, p=1, s=2,
-                        act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.head_elan_4 = ELAN_CSP_Block(in_dim=int(512 * width) + c5,
-                                          out_dim=int(512 * width * ratio),
-                                          expand_ratio=0.5,
-                                          nblocks=int(3*depth),
-                                          shortcut=False,
-                                          depthwise=depthwise,
-                                          norm_type=norm_type,
-                                          act_type=act_type
-                                          )
-
-
-    def forward(self, features):
-        c3, c4, c5 = features
-
-        # Top down
-        ## P5 -> P4
-        c6 = F.interpolate(c5, scale_factor=2.0)
-        c7 = torch.cat([c6, c4], dim=1)
-        c8 = self.head_elan_1(c7)
-        ## P4 -> P3
-        c9 = F.interpolate(c8, scale_factor=2.0)
-        c10 = torch.cat([c9, c3], dim=1)
-        c11 = self.head_elan_2(c10)
-
-        # Bottom up
-        # p3 -> P4
-        c12 = self.mp1(c11)
-        c13 = torch.cat([c12, c8], dim=1)
-        c14 = self.head_elan_3(c13)
-        # P4 -> P5
-        c15 = self.mp2(c14)
-        c16 = torch.cat([c15, c5], dim=1)
-        c17 = self.head_elan_4(c16)
-
-        out_feats = [c11, c14, c17] # [P3, P4, P5]
-        
-        return out_feats
-
-
-def build_fpn(cfg, in_dims):
-    model = cfg['fpn']
-    # build neck
-    if model == 'yolov8_pafpn':
-        fpn_net = Yolov8PaFPN(in_dims=in_dims,
-                             width=cfg['width'],
-                             depth=cfg['depth'],
-                             ratio=cfg['ratio'],
-                             act_type=cfg['fpn_act'],
-                             norm_type=cfg['fpn_norm'],
-                             depthwise=cfg['fpn_depthwise']
-                             )
-    return fpn_net
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'fpn': 'Yolov8PaFPN',
-        'fpn_act': 'silu',
-        'fpn_norm': 'BN',
-        'fpn_depthwise': False,
-        'width': 1.0,
-        'depth': 1.0,
-        'ratio': 1.0,
-    }
-    model = build_fpn(cfg, in_dims=[256, 512, 512])
-    pyramid_feats = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 512, 20, 20)]
-    t0 = time.time()
-    outputs = model(pyramid_feats)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    for out in outputs:
-        print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(pyramid_feats, ), verbose=False)
-    print('==============================')
-    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))

models/detectors/yolox/loss.py

@@ -51,7 +51,7 @@ class Criterion(object):
         return loss_box
 
 
-    def __call__(self, outputs, targets):        
+    def __call__(self, outputs, targets, epoch=0):        
         """
             outputs['pred_obj']: List(Tensor) [B, M, 1]
             outputs['pred_cls']: List(Tensor) [B, M, C]

train.sh

@@ -3,7 +3,7 @@ python train.py \
         --cuda \
         -d coco \
         --root /mnt/share/ssd2/dataset/ \
-        -m yolov5_m \
+        -m yolov7_plus_n \
         -bs 16 \
         -size 640 \
         --wp_epoch 1 \