From 6cfe3643cea6bb1c1e4f8bc81181bdb9b21878de Mon Sep 17 00:00:00 2001
From: George Ni <nemonameless@qq.com>
Date: Tue, 18 May 2021 16:48:56 +0800
Subject: [PATCH] [MOT] fix for_mot in yolo_fpn (#3039)

* fix for_mot in PPYOLOTinyFPN and PPYOLOPAN

* fix yolo_fpn
---
 ppdet/modeling/architectures/yolo.py |  3 +-
 ppdet/modeling/necks/yolo_fpn.py     | 48 +++++++++++++++++++++-------
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/ppdet/modeling/architectures/yolo.py b/ppdet/modeling/architectures/yolo.py
index 276f86e9f..0a035a94a 100644
--- a/ppdet/modeling/architectures/yolo.py
+++ b/ppdet/modeling/architectures/yolo.py
@@ -30,7 +30,8 @@ class YOLOv3(BaseArch):
             yolo_head (nn.Layer): anchor_head instance
             bbox_post_process (object): `BBoxPostProcess` instance
             data_format (str): data format, NCHW or NHWC
-            for_mot (bool): whether return other features used in tracking model 
+            for_mot (bool): whether return other features for multi-object tracking
+                models, default False in pure object detection models.
         """
         super(YOLOv3, self).__init__(data_format=data_format)
         self.backbone = backbone
diff --git a/ppdet/modeling/necks/yolo_fpn.py b/ppdet/modeling/necks/yolo_fpn.py
index 3e65a7db8..1ce653103 100644
--- a/ppdet/modeling/necks/yolo_fpn.py
+++ b/ppdet/modeling/necks/yolo_fpn.py
@@ -18,11 +18,9 @@ import paddle.nn.functional as F
 from paddle import ParamAttr
 from ppdet.core.workspace import register, serializable
 from ..backbones.darknet import ConvBNLayer
-import numpy as np
-
 from ..shape_spec import ShapeSpec
 
-__all__ = ['YOLOv3FPN', 'PPYOLOFPN']
+__all__ = ['YOLOv3FPN', 'PPYOLOFPN', 'PPYOLOTinyFPN', 'PPYOLOPAN']
 
 
 def add_coord(x, data_format):
@@ -492,8 +490,11 @@ class YOLOv3FPN(nn.Layer):
         assert len(blocks) == self.num_blocks
         blocks = blocks[::-1]
         yolo_feats = []
+
+        # add embedding features output for multi-object tracking model
         if for_mot:
             emb_feats = []
+
         for i, block in enumerate(blocks):
             if i > 0:
                 if self.data_format == 'NCHW':
@@ -504,7 +505,7 @@ class YOLOv3FPN(nn.Layer):
             yolo_feats.append(tip)
 
             if for_mot:
-                # add emb_feats output
+                # add embedding features output
                 emb_feats.append(route)
 
             if i < self.num_blocks - 1:
@@ -668,8 +669,11 @@ class PPYOLOFPN(nn.Layer):
         assert len(blocks) == self.num_blocks
         blocks = blocks[::-1]
         yolo_feats = []
+
+        # add embedding features output for multi-object tracking model
         if for_mot:
             emb_feats = []
+
         for i, block in enumerate(blocks):
             if i > 0:
                 if self.data_format == 'NCHW':
@@ -680,7 +684,7 @@ class PPYOLOFPN(nn.Layer):
             yolo_feats.append(tip)
 
             if for_mot:
-                # add emb_feats output
+                # add embedding features output
                 emb_feats.append(route)
 
             if i < self.num_blocks - 1:
@@ -780,11 +784,15 @@ class PPYOLOTinyFPN(nn.Layer):
                         name=name))
                 self.routes.append(route)
 
-    def forward(self, blocks):
+    def forward(self, blocks, for_mot=False):
         assert len(blocks) == self.num_blocks
         blocks = blocks[::-1]
-
         yolo_feats = []
+
+        # add embedding features output for multi-object tracking model
+        if for_mot:
+            emb_feats = []
+
         for i, block in enumerate(blocks):
             if i == 0 and self.spp_:
                 block = self.spp(block)
@@ -797,12 +805,19 @@ class PPYOLOTinyFPN(nn.Layer):
             route, tip = self.yolo_blocks[i](block)
             yolo_feats.append(tip)
 
+            if for_mot:
+                # add embedding features output
+                emb_feats.append(route)
+
             if i < self.num_blocks - 1:
                 route = self.routes[i](route)
                 route = F.interpolate(
                     route, scale_factor=2., data_format=self.data_format)
 
-        return yolo_feats
+        if for_mot:
+            return {'yolo_feats': yolo_feats, 'emb_feats': emb_feats}
+        else:
+            return yolo_feats
 
     @classmethod
     def from_config(cls, cfg, input_shape):
@@ -964,11 +979,15 @@ class PPYOLOPAN(nn.Layer):
 
         self._out_channels = self._out_channels[::-1]
 
-    def forward(self, blocks):
+    def forward(self, blocks, for_mot=False):
         assert len(blocks) == self.num_blocks
         blocks = blocks[::-1]
-        # fpn
         fpn_feats = []
+
+        # add embedding features output for multi-object tracking model
+        if for_mot:
+            emb_feats = []
+
         for i, block in enumerate(blocks):
             if i > 0:
                 if self.data_format == 'NCHW':
@@ -978,6 +997,10 @@ class PPYOLOPAN(nn.Layer):
             route, tip = self.fpn_blocks[i](block)
             fpn_feats.append(tip)
 
+            if for_mot:
+                # add embedding features output
+                emb_feats.append(route)
+
             if i < self.num_blocks - 1:
                 route = self.fpn_routes[i](route)
                 route = F.interpolate(
@@ -996,7 +1019,10 @@ class PPYOLOPAN(nn.Layer):
             route, tip = self.pan_blocks[i](block)
             pan_feats.append(tip)
 
-        return pan_feats[::-1]
+        if for_mot:
+            return {'yolo_feats': pan_feats[::-1], 'emb_feats': emb_feats}
+        else:
+            return pan_feats[::-1]
 
     @classmethod
     def from_config(cls, cfg, input_shape):
-- 
GitLab