|
|
from torch import Tensor |
|
|
|
|
|
from mmdet.models.detectors.single_stage import SingleStageDetector |
|
|
from mmdet.registry import MODELS |
|
|
from mmdet.structures import SampleList |
|
|
from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
|
|
|
|
|
|
|
|
@MODELS.register_module() |
|
|
class XDecoder(SingleStageDetector): |
|
|
|
|
|
def __init__(self, |
|
|
backbone: ConfigType, |
|
|
neck: OptConfigType = None, |
|
|
head: OptConfigType = None, |
|
|
test_cfg: OptConfigType = None, |
|
|
data_preprocessor: OptConfigType = None, |
|
|
init_cfg: OptMultiConfig = None): |
|
|
super(SingleStageDetector, self).__init__( |
|
|
data_preprocessor=data_preprocessor, init_cfg=init_cfg) |
|
|
self.backbone = MODELS.build(backbone) |
|
|
if neck is not None: |
|
|
self.neck = MODELS.build(neck) |
|
|
|
|
|
head_ = head.deepcopy() |
|
|
head_.update(test_cfg=test_cfg) |
|
|
self.sem_seg_head = MODELS.build(head_) |
|
|
|
|
|
def predict(self, |
|
|
batch_inputs: Tensor, |
|
|
batch_data_samples: SampleList, |
|
|
rescale: bool = True) -> SampleList: |
|
|
visual_features = self.extract_feat(batch_inputs) |
|
|
outputs = self.sem_seg_head.predict(visual_features, |
|
|
batch_data_samples) |
|
|
return outputs |
|
|
|
