from __future__ import absolute_import from __future__ import division from __future__ import print_function from ppdet.core.workspace import register, create from .meta_arch import BaseArch __all__ = ['YOLOv3'] @register class YOLOv3(BaseArch): __category__ = 'architecture' __shared__ = ['data_format'] __inject__ = ['post_process'] def __init__(self, backbone='DarkNet', neck='YOLOv3FPN', yolo_head='YOLOv3Head', post_process='BBoxPostProcess', data_format='NCHW'): super(YOLOv3, self).__init__(data_format=data_format) self.backbone = backbone self.neck = neck self.yolo_head = yolo_head self.post_process = post_process @classmethod def from_config(cls, cfg, *args, **kwargs): # backbone backbone = create(cfg['backbone']) # fpn kwargs = {'input_shape': backbone.out_shape} neck = create(cfg['neck'], **kwargs) # head kwargs = {'input_shape': neck.out_shape} yolo_head = create(cfg['yolo_head'], **kwargs) return { 'backbone': backbone, 'neck': neck, "yolo_head": yolo_head, } def _forward(self): body_feats = self.backbone(self.inputs) body_feats = self.neck(body_feats) if self.training: return self.yolo_head(body_feats, self.inputs) else: yolo_head_outs = self.yolo_head(body_feats) bbox, bbox_num = self.post_process( yolo_head_outs, self.yolo_head.mask_anchors, self.inputs['im_shape'], self.inputs['scale_factor']) return bbox, bbox_num def get_loss(self): return self._forward() def get_pred(self): bbox_pred, bbox_num = self._forward() output = {'bbox': bbox_pred, 'bbox_num': bbox_num} return output