Spaces:
Sleeping
Sleeping
File size: 2,153 Bytes
14c9181 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod
from typing import Tuple
from mmengine.model import BaseModule
from torch import Tensor
from mmocr.utils import DetSampleList
class BaseRoIHead(BaseModule, metaclass=ABCMeta):
"""Base class for RoIHeads."""
@property
def with_rec_head(self):
"""bool: whether the RoI head contains a `mask_head`"""
return hasattr(self, 'rec_head') and self.rec_head is not None
@property
def with_extractor(self):
"""bool: whether the RoI head contains a `mask_head`"""
return hasattr(self,
'roi_extractor') and self.roi_extractor is not None
# @abstractmethod
# def init_assigner_sampler(self, *args, **kwargs):
# """Initialize assigner and sampler."""
# pass
@abstractmethod
def loss(self, x: Tuple[Tensor], data_samples: DetSampleList):
"""Perform forward propagation and loss calculation of the roi head on
the features of the upstream network."""
@abstractmethod
def predict(self, x: Tuple[Tensor],
data_samples: DetSampleList) -> DetSampleList:
"""Perform forward propagation of the roi head and predict detection
results on the features of the upstream network.
Args:
x (tuple[Tensor]): Features from upstream network. Each
has shape (N, C, H, W).
data_samples (List[:obj:`DetDataSample`]): The Data
Samples. It usually includes `gt_instance`
Returns:
list[obj:`DetDataSample`]: Detection results of each image.
Each item usually contains following keys in 'pred_instance'
- scores (Tensor): Classification scores, has a shape
(num_instance, )
- labels (Tensor): Labels of bboxes, has a shape
(num_instances, ).
- bboxes (Tensor): Has a shape (num_instances, 4),
the last dimension 4 arrange as (x1, y1, x2, y2).
- polygon (List[Tensor]): Has a shape (num_instances, H, W).
"""
|