fix(box): box_predictor无法导入

Author: zjZSTU

py/ssd/models/box_head/__init__.py

@@ -1 +1 @@
-from .build import build_box_head, build_box_predictor
+from .build import build_box_head

py/ssd/models/box_head/box_head.py

@@ -3,7 +3,7 @@
 
 from ssd.models import registry
 from ssd.models.anchors.prior_box import PriorBox
-from ssd.models.box_head import build_box_predictor
+from ssd.models.box_head.box_predictor import build_box_predictor
 from ssd.utils import box_utils
 from .inference import PostProcessor
 from .loss import MultiBoxLoss

py/ssd/models/box_head/box_predictor.py

@@ -11,7 +11,8 @@ def __init__(self, cfg):
         self.cfg = cfg
         self.cls_headers = nn.ModuleList()
         self.reg_headers = nn.ModuleList()
-        for level, (boxes_per_location, out_channels) in enumerate(zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)):
+        for level, (boxes_per_location, out_channels) in enumerate(
+                zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)):
             self.cls_headers.append(self.cls_block(level, out_channels, boxes_per_location))
             self.reg_headers.append(self.reg_block(level, out_channels, boxes_per_location))
         self.reset_parameters()
@@ -36,7 +37,8 @@ def forward(self, features):
             bbox_pred.append(reg_header(feature).permute(0, 2, 3, 1).contiguous())
 
         batch_size = features[0].shape[0]
-        cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1, self.cfg.MODEL.NUM_CLASSES)
+        cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1,
+                                                                                         self.cfg.MODEL.NUM_CLASSES)
         bbox_pred = torch.cat([l.view(l.shape[0], -1) for l in bbox_pred], dim=1).view(batch_size, -1, 4)
 
         return cls_logits, bbox_pred
@@ -45,7 +47,8 @@ def forward(self, features):
 @registry.BOX_PREDICTORS.register('SSDBoxPredictor')
 class SSDBoxPredictor(BoxPredictor):
     def cls_block(self, level, out_channels, boxes_per_location):
-        return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+        return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1,
+                         padding=1)
 
     def reg_block(self, level, out_channels, boxes_per_location):
         return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
@@ -57,10 +60,15 @@ def cls_block(self, level, out_channels, boxes_per_location):
         num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
         if level == num_levels - 1:
             return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=1)
-        return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+        return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1,
+                               padding=1)
 
     def reg_block(self, level, out_channels, boxes_per_location):
         num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
         if level == num_levels - 1:
             return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=1)
-        return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+        return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+def build_box_predictor(cfg):
+    return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg)

py/ssd/models/box_head/build.py

@@ -9,12 +9,7 @@
 
 from ssd.models import registry
 from .box_head import SSDBoxHead
-from .box_predictor import SSDBoxPredictor, SSDLiteBoxPredictor
 
 
 def build_box_head(cfg):
     return registry.BOX_HEADS[cfg.MODEL.BOX_HEAD.NAME](cfg)
-
-
-def build_box_predictor(cfg):
-    return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg)