更改mask2former接口，适配maas

        Link: https://code.alibaba-inc.com/pai-vision/EasyCV/codereview/10292532

* modify postprocess

* Merge remote-tracking branch 'remotes/origin/master' into mask2former_postprocessing

# Conflicts:
#	tests/ut_config.py

* assert the output value

* add mask2former models to data/test/xxx/models

* fixed train

Author: tianchu.gtc

configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py

@@ -15,6 +15,26 @@
     'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
     'hair drier', 'toothbrush'
 ]
+PALETTE = [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230),
+           (106, 0, 228), (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192),
+           (250, 170, 30), (100, 170, 30), (220, 220, 0), (175, 116, 175),
+           (250, 0, 30), (165, 42, 42), (255, 77, 255), (0, 226, 252),
+           (182, 182, 255), (0, 82, 0), (120, 166, 157), (110, 76, 0),
+           (174, 57, 255), (199, 100, 0), (72, 0, 118), (255, 179, 240),
+           (0, 125, 92), (209, 0, 151), (188, 208, 182), (0, 220, 176),
+           (255, 99, 164), (92, 0, 73), (133, 129, 255), (78, 180, 255),
+           (0, 228, 0), (174, 255, 243), (45, 89, 255), (134, 134, 103),
+           (145, 148, 174), (255, 208, 186), (197, 226, 255), (171, 134, 1),
+           (109, 63, 54), (207, 138, 255), (151, 0, 95), (9, 80, 61),
+           (84, 105, 51), (74, 65, 105), (166, 196, 102), (208, 195, 210),
+           (255, 109, 65), (0, 143, 149), (179, 0, 194), (209, 99, 106),
+           (5, 121, 0), (227, 255, 205), (147, 186, 208), (153, 69, 1),
+           (3, 95, 161), (163, 255, 0), (119, 0, 170), (0, 182, 199),
+           (0, 165, 120), (183, 130, 88), (95, 32, 0), (130, 114, 135),
+           (110, 129, 133), (166, 74, 118), (219, 142, 185), (79, 210, 114),
+           (178, 90, 62), (65, 70, 15), (127, 167, 115), (59, 105, 106),
+           (142, 108, 45), (196, 172, 0), (95, 54, 80), (128, 76, 255),
+           (201, 57, 1), (246, 0, 122), (191, 162, 208)]
 
 model = dict(
     type='Mask2Former',

data/test/segmentation/models/mask2former_pan_export.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:163a344e29b965cdb6c6c24e189e84a269580d63237253f359de35e944ec5421
+size 528712836

data/test/segmentation/models/mask2former_r50_instance.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a33e6b1f5623057c6920226767c91a44a072acc27ece5ba24fdeb2a9a1bb2ba2
+size 528548036

easycv/models/segmentation/mask2former.py

@@ -110,8 +110,13 @@ def init_weights(self):
             print_log('load model from init weights')
             self.backbone.init_weights()
 
-    def forward_train(self, img, gt_labels, gt_masks, gt_semantic_seg,
-                      img_metas):
+    def forward_train(self,
+                      img,
+                      gt_labels,
+                      gt_masks=None,
+                      gt_semantic_seg=None,
+                      img_metas=None,
+                      **kwargs):
         features = self.backbone(img)
         outputs = self.head(features)
         targets = self.preprocess_gt(gt_labels, gt_masks, gt_semantic_seg,
@@ -125,7 +130,12 @@ def forward_train(self, img, gt_labels, gt_masks, gt_semantic_seg,
                 losses.pop(k)
         return losses
 
-    def forward_test(self, img, img_metas, rescale=True, encode=True):
+    def forward_test(self,
+                     img,
+                     img_metas,
+                     rescale=True,
+                     encode=True,
+                     **kwargs):
         features = self.backbone(img[0])
         outputs = self.head(features)
         mask_cls_results = outputs['pred_logits']
@@ -189,23 +199,6 @@ def forward_test(self, img, img_metas, rescale=True, encode=True):
         outputs['pan_results'] = pan_masks
         return outputs
 
-    def forward(self,
-                img,
-                mode='train',
-                gt_labels=None,
-                gt_masks=None,
-                gt_semantic_seg=None,
-                img_metas=None,
-                **kwargs):
-
-        if mode == 'train':
-            return self.forward_train(img, gt_labels, gt_masks,
-                                      gt_semantic_seg, img_metas)
-        elif mode == 'test':
-            return self.forward_test(img, img_metas)
-        else:
-            raise Exception('No such mode: {}'.format(mode))
-
     def instance_postprocess(self, mask_cls, mask_pred):
         """Instance segmengation postprocess.
 
@@ -233,8 +226,10 @@ def instance_postprocess(self, mask_cls, mask_pred):
         # shape (num_queries, num_class)
         scores = F.softmax(mask_cls, dim=-1)[:, :-1]
         # shape (num_queries * num_class, )
-        labels = torch.arange(self.num_classes, device=mask_cls.device).\
-            unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
+        labels = torch.arange(
+            self.num_classes,
+            device=mask_cls.device).unsqueeze(0).repeat(num_queries,
+                                                        1).flatten(0, 1)
         scores_per_image, top_indices = scores.flatten(0, 1).topk(
             max_per_image, sorted=False)
         labels_per_image = labels[top_indices]

easycv/predictors/segmentation.py

@@ -167,29 +167,44 @@ def forward(self, inputs):
         """Model forward.
         """
         with torch.no_grad():
-            outputs = self.model(**inputs, mode='test', encode=False)
+            outputs = self.model.forward(**inputs, mode='test', encode=False)
+        if self.task_mode == 'instance':
+            outputs.pop('pan_results')
+        elif self.task_mode == 'panoptic':
+            outputs.pop('detection_masks')
+            outputs.pop('detection_boxes')
+            outputs.pop('detection_scores')
+            outputs.pop('detection_classes')
         return outputs
 
-    def postprocess(self, inputs):
+    def postprocess_single(self, inputs, *args, **kwargs):
         output = {}
         if self.task_mode == 'panoptic':
-            output['pan'] = inputs['pan_results'][0]
+            pan_results = inputs['pan_results']
+            # keep objects ahead
+            ids = np.unique(pan_results)[::-1]
+            legal_indices = ids != len(self.CLASSES)  # for VOID label
+            ids = ids[legal_indices]
+            labels = np.array([id % 1000 for id in ids], dtype=np.int64)
+            segms = (pan_results[None] == ids[:, None, None])
+            masks = [it.astype(np.int) for it in segms]
+            labels_txt = np.array(self.CLASSES)[labels].tolist()
+
+            output['masks'] = masks
+            output['labels'] = labels_txt
+            output['labels_ids'] = labels
         elif self.task_mode == 'instance':
-            output['segms'] = inputs['detection_masks'][0]
-            output['bboxes'] = inputs['detection_boxes'][0]
-            output['scores'] = inputs['detection_scores'][0]
-            output['labels'] = inputs['detection_classes'][0]
+            output['segms'] = inputs['detection_masks']
+            output['bboxes'] = inputs['detection_boxes']
+            output['scores'] = inputs['detection_scores']
+            output['labels'] = inputs['detection_classes']
         else:
             raise ValueError(f'Not support model {self.task_mode}')
         return output
 
-    def show_panoptic(self, img, pan_mask):
-        pan_label = np.unique(pan_mask)
-        pan_label = pan_label[pan_label % 1000 != self.classes]
-        masks = np.array([pan_mask == num for num in pan_label])
-
+    def show_panoptic(self, img, masks, labels):
         palette = np.asarray(self.cfg.PALETTE)
-        palette = palette[pan_label % 1000]
+        palette = palette[labels % 1000]
         panoptic_result = draw_masks(img, masks, palette)
         return panoptic_result
 
@@ -199,10 +214,11 @@ def show_instance(self, img, segms, bboxes, scores, labels, score_thr=0.5):
             bboxes = bboxes[inds, :]
             segms = segms[inds, ...]
             labels = labels[inds]
-        palette = np.asarray(self.cfg.PALETTE)
+        palette = np.asarray(self.PALETTE)
         palette = palette[labels]
+
         instance_result = draw_masks(img, segms, palette)
-        class_name = np.array(self.class_name)
+        class_name = np.array(self.CLASSES)
         instance_result = imshow_bboxes(
             instance_result, bboxes, class_name[labels], show=False)
         return instance_result

tests/predictors/test_segmentation.py

@@ -5,15 +5,19 @@
 import tempfile
 import unittest
 
+import cv2
 import numpy as np
 from mmcv import Config
 from PIL import Image
-from tests.ut_config import (MODEL_CONFIG_SEGFORMER,
+from tests.ut_config import (MODEL_CONFIG_MASK2FORMER_INS,
+                             MODEL_CONFIG_MASK2FORMER_PAN,
+                             MODEL_CONFIG_SEGFORMER,
                              PRETRAINED_MODEL_MASK2FORMER_DIR,
                              PRETRAINED_MODEL_SEGFORMER, TEST_IMAGES_DIR)
 
 from easycv.file import io
-from easycv.predictors.segmentation import SegmentationPredictor
+from easycv.predictors.segmentation import (Mask2formerPredictor,
+                                            SegmentationPredictor)
 
 
 class SegmentationPredictorTest(unittest.TestCase):
@@ -112,34 +116,94 @@ def test_dump(self):
         shutil.rmtree(temp_dir, ignore_errors=True)
 
 
-@unittest.skipIf(True, 'WIP')
 class Mask2formerPredictorTest(unittest.TestCase):
 
-    def test_single(self):
-        import cv2
-        from easycv.predictors.segmentation import Mask2formerPredictor
-        pan_ckpt = os.path.join(PRETRAINED_MODEL_MASK2FORMER_DIR,
-                                'mask2former_pan_export.pth')
-        instance_ckpt = os.path.join(PRETRAINED_MODEL_MASK2FORMER_DIR,
-                                     'mask2former_r50_instance.pth')
-        img_path = os.path.join(TEST_IMAGES_DIR, 'mask2former.jpg')
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        self.img_path = './data/test/segmentation/data/000000309022.jpg'
+        self.pan_ckpt = './data/test/segmentation/models/mask2former_pan_export.pth'
+        self.instance_ckpt = './data/test/segmentation/models/mask2former_r50_instance.pth'
 
+    def test_panoptic_single(self):
         # panop
+        segmentation_model_config = MODEL_CONFIG_MASK2FORMER_PAN
         predictor = Mask2formerPredictor(
-            model_path=pan_ckpt, output_mode='panoptic')
-        img = cv2.imread(img_path)
-        predict_out = predictor([img])
-        pan_img = predictor.show_panoptic(img, predict_out[0]['pan'])
+            model_path=self.pan_ckpt,
+            task_mode='panoptic',
+            config_file=segmentation_model_config)
+        img = cv2.imread(self.img_path)
+        predict_out = predictor([self.img_path])
+        self.assertEqual(len(predict_out), 1)
+        self.assertEqual(len(predict_out[0]['masks']), 14)
+        self.assertListEqual(
+            predict_out[0]['labels_ids'].tolist(),
+            [71, 69, 39, 39, 39, 128, 127, 122, 118, 115, 111, 104, 84, 83])
+
+        pan_img = predictor.show_panoptic(
+            img,
+            masks=predict_out[0]['masks'],
+            labels=predict_out[0]['labels_ids'])
         cv2.imwrite('pan_out.jpg', pan_img)
 
+    def test_panoptic_batch(self):
+        total_samples = 2
+        segmentation_model_config = MODEL_CONFIG_MASK2FORMER_PAN
+        predictor = Mask2formerPredictor(
+            model_path=self.pan_ckpt,
+            task_mode='panoptic',
+            config_file=segmentation_model_config,
+            batch_size=total_samples)
+        predict_out = predictor([self.img_path] * total_samples)
+        self.assertEqual(len(predict_out), total_samples)
+        img = cv2.imread(self.img_path)
+        for i in range(total_samples):
+            save_name = 'pan_out_batch_%d.jpg' % i
+            self.assertEqual(len(predict_out[i]['masks']), 14)
+            self.assertListEqual(predict_out[i]['labels_ids'].tolist(), [
+                71, 69, 39, 39, 39, 128, 127, 122, 118, 115, 111, 104, 84, 83
+            ])
+            pan_img = predictor.show_panoptic(
+                img,
+                masks=predict_out[i]['masks'],
+                labels=predict_out[i]['labels_ids'])
+            cv2.imwrite(save_name, pan_img)
+
+    def test_instance_single(self):
         # instance
+        segmentation_model_config = MODEL_CONFIG_MASK2FORMER_INS
         predictor = Mask2formerPredictor(
-            model_path=instance_ckpt, output_mode='instance')
-        img = cv2.imread(img_path)
-        predict_out = predictor.predict([img], mode='instance')
+            model_path=self.instance_ckpt,
+            task_mode='instance',
+            config_file=segmentation_model_config)
+        img = cv2.imread(self.img_path)
+        predict_out = predictor([self.img_path])
+        self.assertEqual(len(predict_out), 1)
+        self.assertEqual(predict_out[0]['segms'].shape, (100, 480, 640))
+        self.assertListEqual(predict_out[0]['labels'][:10].tolist(),
+                             [41, 69, 72, 45, 68, 70, 41, 69, 69, 45])
+
         instance_img = predictor.show_instance(img, **predict_out[0])
         cv2.imwrite('instance_out.jpg', instance_img)
 
+    def test_instance_batch(self):
+        total_samples = 2
+        segmentation_model_config = MODEL_CONFIG_MASK2FORMER_INS
+        predictor = Mask2formerPredictor(
+            model_path=self.instance_ckpt,
+            task_mode='instance',
+            config_file=segmentation_model_config,
+            batch_size=total_samples)
+        img = cv2.imread(self.img_path)
+        predict_out = predictor([self.img_path] * total_samples)
+        self.assertEqual(len(predict_out), total_samples)
+        for i in range(total_samples):
+            save_name = 'instance_out_batch_%d.jpg' % i
+            self.assertEqual(predict_out[i]['segms'].shape, (100, 480, 640))
+            self.assertListEqual(predict_out[0]['labels'][:10].tolist(),
+                                 [41, 69, 72, 45, 68, 70, 41, 69, 69, 45])
+            instance_img = predictor.show_instance(img, **(predict_out[i]))
+            cv2.imwrite(save_name, instance_img)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/ut_config.py

@@ -132,3 +132,7 @@
     './configs/segmentation/segformer/segformer_b0_coco.py')
 SMALL_COCO_WHOLE_BODY_HAND_ROOT = 'data/test/pose/hand/small_whole_body_hand_coco'
 SMALL_COCO_WHOLEBODY_ROOT = 'data/test/pose/wholebody/data'
+MODEL_CONFIG_MASK2FORMER_PAN = (
+    './configs/segmentation/mask2former/mask2former_r50_8xb2_e50_panoptic.py')
+MODEL_CONFIG_MASK2FORMER_INS = (
+    './configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py')