[to #42322933] add abnormal detection models

添加了针对长尾/小目标问题解决的异常视觉检测模型

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11536379

modelscope/metainfo.py

@@ -15,6 +15,7 @@ class Models(object):
     tinynas_damoyolo = 'tinynas-damoyolo'
     # vision models
     detection = 'detection'
+    mask_scoring = 'MaskScoring'
     image_restoration = 'image-restoration'
     realtime_object_detection = 'realtime-object-detection'
     realtime_video_object_detection = 'realtime-video-object-detection'
@@ -233,6 +234,7 @@ class Pipelines(object):
     hand_2d_keypoints = 'hrnetv2w18_hand-2d-keypoints_image'
     human_detection = 'resnet18-human-detection'
     object_detection = 'vit-object-detection'
+    abnormal_object_detection = 'abnormal-object-detection'
     easycv_detection = 'easycv-detection'
     easycv_segmentation = 'easycv-segmentation'
     face_2d_keypoints = 'mobilenet_face-2d-keypoints_alignment'

modelscope/models/cv/abnormal_object_detection/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    from .mmdet_model import AbnormalDetectionModel
+
+else:
+    _import_structure = {'mmdet_model': ['AbnormalDetectionModel']}
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/cv/abnormal_object_detection/mmdet_model.py

@@ -0,0 +1,103 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path as osp
+
+import numpy as np
+import torch
+
+from modelscope.metainfo import Models
+from modelscope.models.base.base_torch_model import TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.utils.config import Config
+from modelscope.utils.constant import ModelFile, Tasks
+from .mmdet_ms import MaskScoringNRoIHead, SingleRoINExtractor
+
+
+@MODELS.register_module(
+    Tasks.image_object_detection, module_name=Models.mask_scoring)
+class AbnormalDetectionModel(TorchModel):
+
+    def __init__(self, model_dir: str, *args, **kwargs):
+        """str -- model file root."""
+        super().__init__(model_dir, *args, **kwargs)
+
+        from mmcv.runner import load_checkpoint
+        from mmdet.datasets import replace_ImageToTensor
+        from mmdet.datasets.pipelines import Compose
+        from mmdet.models import build_detector
+
+        model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE)
+        config_path = osp.join(model_dir, 'mmcv_config.py')
+        config = Config.from_file(config_path)
+        config.model.pretrained = None
+        self.model = build_detector(
+            config.model, test_cfg=config.get('test_cfg'))
+
+        checkpoint = load_checkpoint(
+            self.model, model_path, map_location='cpu')
+        self.class_names = checkpoint['meta']['CLASSES']
+        config.test_pipeline[0].type = 'LoadImageFromWebcam'
+        self.transform_input = Compose(
+            replace_ImageToTensor(config.test_pipeline))
+        self.model.cfg = config
+        self.model.eval()
+        self.score_thr = config.score_thr
+
+    def inference(self, data):
+        """data is dict,contain img and img_metas,follow with mmdet.
+        Args:
+            imgs (List[Tensor]): the outer list indicates test-time
+                augmentations and inner Tensor should have a shape NxCxHxW,
+                which contains all images in the batch.
+            img_metas (List[List[dict]]): the outer list indicates test-time
+                augs (multiscale, flip, etc.) and the inner list indicates
+                images in a batch.
+        """
+
+        with torch.no_grad():
+            results = self.model(
+                return_loss=False,
+                rescale=True,
+                img=data['img'],
+                img_metas=data['img_metas'])
+        return results
+
+    def preprocess(self, image):
+        """image is numpy return is dict contain img and img_metas,follow with mmdet."""
+
+        from mmcv.parallel import collate, scatter
+        data = dict(img=image)
+        data = self.transform_input(data)
+        data = collate([data], samples_per_gpu=1)
+        data['img_metas'] = [
+            img_metas.data[0] for img_metas in data['img_metas']
+        ]
+        data['img'] = [img.data[0] for img in data['img']]
+
+        if next(self.model.parameters()).is_cuda:
+            data = scatter(data, [next(self.model.parameters()).device])[0]
+
+        return data
+
+    def postprocess(self, inputs):
+
+        if isinstance(inputs[0], tuple):
+            bbox_result, _ = inputs[0]
+        else:
+            bbox_result, _ = inputs[0], None
+        labels = [
+            np.full(bbox.shape[0], i, dtype=np.int32)
+            for i, bbox in enumerate(bbox_result)
+        ]
+        labels = np.concatenate(labels)
+
+        bbox_result = np.vstack(bbox_result)
+        scores = bbox_result[:, -1]
+        inds = scores > self.score_thr
+        if np.sum(np.array(inds).astype('int')) == 0:
+            return None, None, None
+        bboxes = bbox_result[inds, :]
+        labels = labels[inds]
+        scores = np.around(bboxes[:, 4], 6)
+        bboxes = (bboxes[:, 0:4]).astype(int)
+        labels = [self.class_names[i_label] for i_label in labels]
+        return bboxes, scores, labels

modelscope/models/cv/abnormal_object_detection/mmdet_ms/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .roi_head import MaskScoringNRoIHead, SingleRoINExtractor

modelscope/models/cv/abnormal_object_detection/mmdet_ms/roi_head/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .mask_scoring_roi_head import MaskScoringNRoIHead
+from .roi_extractors import SingleRoINExtractor
+
+__all__ = ['MaskScoringNRoIHead', 'SingleRoINExtractor']

modelscope/models/cv/abnormal_object_detection/mmdet_ms/roi_head/mask_scoring_roi_head.py

@@ -0,0 +1,138 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# Implementation in this file is modified based on mmdetection
+# Originally Apache 2.0 License and publicly avaialbe at https://github.com/open-mmlab/mmdetection
+import torch
+from mmdet.core import bbox2roi
+from mmdet.models.builder import HEADS, build_head
+from mmdet.models.roi_heads.standard_roi_head import StandardRoIHead
+
+
+@HEADS.register_module()
+class MaskScoringNRoIHead(StandardRoIHead):
+    """Mask Scoring RoIHead for Mask Scoring RCNN.
+
+    https://arxiv.org/abs/1903.00241
+    """
+
+    def __init__(self,
+                 mask_iou_head=None,
+                 reg_roi_scale_factor=None,
+                 **kwargs):
+        # assert mask_iou_head is not None
+        super(MaskScoringNRoIHead, self).__init__(**kwargs)
+        if mask_iou_head is not None:
+            self.mask_iou_head = build_head(mask_iou_head)
+        self.reg_roi_scale_factor = reg_roi_scale_factor
+
+    def _bbox_forward(self, x, rois):
+        """Box head forward function used in both training and testing time."""
+        bbox_cls_feats = self.bbox_roi_extractor(
+            x[:self.bbox_roi_extractor.num_inputs], rois)
+        bbox_reg_feats = self.bbox_roi_extractor(
+            x[:self.bbox_roi_extractor.num_inputs],
+            rois,
+            roi_scale_factor=self.reg_roi_scale_factor)
+        if self.with_shared_head:
+            bbox_cls_feats = self.shared_head(bbox_cls_feats)
+            bbox_reg_feats = self.shared_head(bbox_reg_feats)
+        cls_score, bbox_pred = self.bbox_head(bbox_cls_feats, bbox_reg_feats)
+
+        bbox_results = dict(
+            cls_score=cls_score,
+            bbox_pred=bbox_pred,
+            bbox_feats=bbox_cls_feats)
+        return bbox_results
+
+    def _mask_forward_train(self, x, sampling_results, bbox_feats, gt_masks,
+                            img_metas):
+        """Run forward function and calculate loss for Mask head in
+        training."""
+        pos_labels = torch.cat([res.pos_gt_labels for res in sampling_results])
+        mask_results = super(MaskScoringNRoIHead,
+                             self)._mask_forward_train(x, sampling_results,
+                                                       bbox_feats, gt_masks,
+                                                       img_metas)
+        if mask_results['loss_mask'] is None:
+            return mask_results
+
+        # mask iou head forward and loss
+        pos_mask_pred = mask_results['mask_pred'][
+            range(mask_results['mask_pred'].size(0)), pos_labels]
+        mask_iou_pred = self.mask_iou_head(mask_results['mask_feats'],
+                                           pos_mask_pred)
+        pos_mask_iou_pred = mask_iou_pred[range(mask_iou_pred.size(0)),
+                                          pos_labels]
+
+        mask_iou_targets = self.mask_iou_head.get_targets(
+            sampling_results, gt_masks, pos_mask_pred,
+            mask_results['mask_targets'], self.train_cfg)
+        loss_mask_iou = self.mask_iou_head.loss(pos_mask_iou_pred,
+                                                mask_iou_targets)
+        mask_results['loss_mask'].update(loss_mask_iou)
+        return mask_results
+
+    def simple_test_mask(self,
+                         x,
+                         img_metas,
+                         det_bboxes,
+                         det_labels,
+                         rescale=False):
+        """Obtain mask prediction without augmentation."""
+        # image shapes of images in the batch
+        ori_shapes = tuple(meta['ori_shape'] for meta in img_metas)
+        scale_factors = tuple(meta['scale_factor'] for meta in img_metas)
+
+        num_imgs = len(det_bboxes)
+        if all(det_bbox.shape[0] == 0 for det_bbox in det_bboxes):
+            num_classes = self.mask_head.num_classes
+            segm_results = [[[] for _ in range(num_classes)]
+                            for _ in range(num_imgs)]
+            mask_scores = [[[] for _ in range(num_classes)]
+                           for _ in range(num_imgs)]
+        else:
+            # if det_bboxes is rescaled to the original image size, we need to
+            # rescale it back to the testing scale to obtain RoIs.
+            if rescale and not isinstance(scale_factors[0], float):
+                scale_factors = [
+                    torch.from_numpy(scale_factor).to(det_bboxes[0].device)
+                    for scale_factor in scale_factors
+                ]
+            _bboxes = [
+                det_bboxes[i][:, :4]
+                * scale_factors[i] if rescale else det_bboxes[i]
+                for i in range(num_imgs)
+            ]
+            mask_rois = bbox2roi(_bboxes)
+            mask_results = self._mask_forward(x, mask_rois)
+            concat_det_labels = torch.cat(det_labels)
+            # get mask scores with mask iou head
+            mask_feats = mask_results['mask_feats']
+            mask_pred = mask_results['mask_pred']
+            mask_iou_pred = self.mask_iou_head(
+                mask_feats, mask_pred[range(concat_det_labels.size(0)),
+                                      concat_det_labels])
+            # split batch mask prediction back to each image
+            num_bboxes_per_img = tuple(len(_bbox) for _bbox in _bboxes)
+            mask_preds = mask_pred.split(num_bboxes_per_img, 0)
+            mask_iou_preds = mask_iou_pred.split(num_bboxes_per_img, 0)
+
+            # apply mask post-processing to each image individually
+            segm_results = []
+            mask_scores = []
+            for i in range(num_imgs):
+                if det_bboxes[i].shape[0] == 0:
+                    segm_results.append(
+                        [[] for _ in range(self.mask_head.num_classes)])
+                    mask_scores.append(
+                        [[] for _ in range(self.mask_head.num_classes)])
+                else:
+                    segm_result = self.mask_head.get_seg_masks(
+                        mask_preds[i], _bboxes[i], det_labels[i],
+                        self.test_cfg, ori_shapes[i], scale_factors[i],
+                        rescale)
+                    # get mask scores with mask iou head
+                    mask_score = self.mask_iou_head.get_mask_scores(
+                        mask_iou_preds[i], det_bboxes[i], det_labels[i])
+                    segm_results.append(segm_result)
+                    mask_scores.append(mask_score)
+        return list(zip(segm_results, mask_scores))

modelscope/models/cv/abnormal_object_detection/mmdet_ms/roi_head/roi_extractors/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .single_level_roi_extractor import SingleRoINExtractor
+
+__all__ = ['SingleRoINExtractor']

modelscope/models/cv/abnormal_object_detection/mmdet_ms/roi_head/roi_extractors/single_level_roi_extractor.py

@@ -0,0 +1,153 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# Implementation in this file is modified based on mmdetection
+# Originally Apache 2.0 License and publicly avaialbe at https://github.com/open-mmlab/mmdetection
+import torch
+from mmcv.runner import force_fp32
+from mmdet.models.builder import ROI_EXTRACTORS
+from mmdet.models.roi_heads.roi_extractors.base_roi_extractor import \
+    BaseRoIExtractor
+
+
+@ROI_EXTRACTORS.register_module()
+class SingleRoINExtractor(BaseRoIExtractor):
+    """Extract RoI features from a single level feature map.
+
+    If there are multiple input feature levels, each RoI is mapped to a level
+    according to its scale. The mapping rule is proposed in
+    `FPN <https://arxiv.org/abs/1612.03144>`_.
+
+    Args:
+        roi_layer (dict): Specify RoI layer type and arguments.
+        out_channels (int): Output channels of RoI layers.
+        featmap_strides (List[int]): Strides of input feature maps.
+        finest_scale (int): Scale threshold of mapping to level 0. Default: 56.
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+            Default: None
+    """
+
+    def __init__(self,
+                 roi_layer,
+                 out_channels,
+                 featmap_strides,
+                 finest_scale=56,
+                 init_cfg=None,
+                 gc_context=False,
+                 offset_feature=False):
+        super(SingleRoINExtractor, self).__init__(roi_layer, out_channels,
+                                                  featmap_strides, init_cfg)
+        self.finest_scale = finest_scale
+        self.gc_context = gc_context
+        self.offset_feature = offset_feature
+        self.pool = torch.nn.AdaptiveAvgPool2d(7)
+
+    def map_roi_levels(self, rois, num_levels):
+        """Map rois to corresponding feature levels by scales.
+
+        - scale < finest_scale * 2: level 0
+        - finest_scale * 2 <= scale < finest_scale * 4: level 1
+        - finest_scale * 4 <= scale < finest_scale * 8: level 2
+        - scale >= finest_scale * 8: level 3
+
+        Args:
+            rois (Tensor): Input RoIs, shape (k, 5).
+            num_levels (int): Total level number.
+
+        Returns:
+            Tensor: Level index (0-based) of each RoI, shape (k, )
+        """
+        a = rois[:, 3] - rois[:, 1]
+        b = rois[:, 4] - rois[:, 2]
+        scale = torch.sqrt(a * b)
+        target_lvls = torch.floor(torch.log2(scale / self.finest_scale + 1e-6))
+        target_lvls = target_lvls.clamp(min=0, max=num_levels - 1).long()
+        return target_lvls
+
+    @force_fp32(apply_to=('feats', ), out_fp16=True)
+    def forward(self, feats, rois, roi_scale_factor=None):
+        """Forward function."""
+        out_size = self.roi_layers[0].output_size
+        num_levels = len(feats)
+        expand_dims = (-1, self.out_channels * out_size[0] * out_size[1])
+        if torch.onnx.is_in_onnx_export():
+            # Work around to export mask-rcnn to onnx
+            roi_feats = rois[:, :1].clone().detach()
+            roi_feats = roi_feats.expand(*expand_dims)
+            roi_feats = roi_feats.reshape(-1, self.out_channels, *out_size)
+            roi_feats = roi_feats * 0
+        else:
+            roi_feats = feats[0].new_zeros(
+                rois.size(0), self.out_channels, *out_size)
+        # TODO: remove this when parrots supports
+        if torch.__version__ == 'parrots':
+            roi_feats.requires_grad = True
+
+        if num_levels == 1:
+            if len(rois) == 0:
+                return roi_feats
+            return self.roi_layers[0](feats[0], rois)
+
+        if self.gc_context:
+            context = []
+            for feat in feats:
+                context.append(self.pool(feat))
+
+        batch_size = feats[0].shape[0]
+        target_lvls = self.map_roi_levels(rois, num_levels)
+
+        if roi_scale_factor is not None:
+            rois = self.roi_rescale(rois, roi_scale_factor)
+
+        for i in range(num_levels):
+            mask = target_lvls == i
+            if torch.onnx.is_in_onnx_export():
+                # To keep all roi_align nodes exported to onnx
+                # and skip nonzero op
+                mask = mask.float().unsqueeze(-1)
+                # select target level rois and reset the rest rois to zero.
+                rois_i = rois.clone().detach()
+                rois_i *= mask
+                mask_exp = mask.expand(*expand_dims).reshape(roi_feats.shape)
+                roi_feats_t = self.roi_layers[i](feats[i], rois_i)
+                roi_feats_t *= mask_exp
+                roi_feats += roi_feats_t
+                continue
+            inds = mask.nonzero(as_tuple=False).squeeze(1)
+            if inds.numel() > 0:
+                rois_ = rois[inds]
+                # todo offset
+                rois_offset = rois[inds]
+                offset = torch.zeros(rois_.size(0), 5)
+                _, _, x_max, y_max = rois_[:, 1].min().item(), rois_[:, 2].min(
+                ).item(), rois_[:, 3].max().item(), rois_[:, 4].max().item()
+                offset[:, 1:3] = -100 * torch.ones(rois_.size(0), 1)
+                offset[:, 3:5] = 100 * torch.ones(rois_.size(0), 1)
+                rois_offset += offset.cuda()
+                rois_offset_thsxy = torch.clamp(rois_offset[:, 1:3], min=0.)
+                rois_offset_ths_xmax = torch.clamp(
+                    rois_offset[:, 3], max=x_max)
+                rois_offset_ths_ymax = torch.clamp(
+                    rois_offset[:, 4], max=y_max)
+                rois_offset[:, 1:3] = rois_offset_thsxy
+                rois_offset[:,
+                            3], rois_offset[:,
+                                            4] = rois_offset_ths_xmax, rois_offset_ths_ymax
+                roi_feats_t = self.roi_layers[i](feats[i], rois_)
+                roi_feats_t_offset = self.roi_layers[i](feats[i], rois_offset)
+                if self.gc_context:
+                    for j in range(batch_size):
+                        roi_feats_t[rois_[:, 0] == j] += context[i][j]
+                elif self.offset_feature:
+                    roi_feats_t += roi_feats_t_offset
+
+                roi_feats[inds] = roi_feats_t
+            else:
+                # Sometimes some pyramid levels will not be used for RoI
+                # feature extraction and this will cause an incomplete
+                # computation graph in one GPU, which is different from those
+                # in other GPUs and will cause a hanging error.
+                # Therefore, we add it to ensure each feature pyramid is
+                # included in the computation graph to avoid runtime bugs.
+                roi_feats += sum(
+                    x.view(-1)[0]
+                    for x in self.parameters()) * 0. + feats[i].sum() * 0.
+        return roi_feats

modelscope/pipelines/cv/image_detection_pipeline.py

@@ -10,13 +10,15 @@ from modelscope.pipelines.base import Input, Pipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.preprocessors import LoadImage
 from modelscope.utils.constant import Tasks
-from modelscope.utils.logger import get_logger
 
 
 @PIPELINES.register_module(
     Tasks.human_detection, module_name=Pipelines.human_detection)
 @PIPELINES.register_module(
     Tasks.image_object_detection, module_name=Pipelines.object_detection)
+@PIPELINES.register_module(
+    Tasks.image_object_detection,
+    module_name=Pipelines.abnormal_object_detection)
 class ImageDetectionPipeline(Pipeline):
 
     def __init__(self, model: str, **kwargs):

tests/pipelines/test_abnormal_object_detection.py

@@ -0,0 +1,29 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+
+class ObjectDetectionTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.image_object_detection
+        self.model_id = 'damo/cv_resnet50_object-detection_maskscoring'
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_abnormal_object_detection(self):
+        input_location = 'data/test/images/image_detection.jpg'
+        object_detect = pipeline(self.task, model=self.model_id)
+        result = object_detect(input_location)
+        print(result)
+
+    @unittest.skip('demo compatibility test is only enabled on a needed-basis')
+    def test_demo_compatibility(self):
+        self.compatibility_check()
+
+
+if __name__ == '__main__':
+    unittest.main()