[Fix] Fix batch inference error for Mask R-CNN (#1575)

* fix Mask R-CNN for multi-batch

* fix flake8

* fix bug

* use Sequence instead list

* fix docstring

* only test_img accept list input

Author: hanrui1sensetime

mmdeploy/apis/visualize.py

@@ -11,7 +11,7 @@
 def visualize_model(model_cfg: Union[str, mmcv.Config],
                     deploy_cfg: Union[str, mmcv.Config],
                     model: Union[str, Sequence[str]],
-                    img: Union[str, np.ndarray],
+                    img: Union[str, np.ndarray, Sequence[str]],
                     device: str,
                     backend: Optional[Backend] = None,
                     output_file: Optional[str] = None,
@@ -36,7 +36,8 @@ def visualize_model(model_cfg: Union[str, mmcv.Config],
         deploy_cfg (str | mmcv.Config): Deployment config file or Config
             object.
         model (str | list[str], BaseSubtask): Input model or file(s).
-        img (str | np.ndarray): Input image file or numpy array for inference.
+        img (str | np.ndarray | Sequence[str]): Input image file(s) or numpy
+            array for inference.
         device (str): A string specifying device type.
         backend (Backend): Specifying backend type, defaults to `None`.
         output_file (str): Output file to save visualized image, defaults to
@@ -74,14 +75,16 @@ def visualize_model(model_cfg: Union[str, mmcv.Config],
         # check headless
         import tkinter
         tkinter.Tk()
-
-        task_processor.visualize(
-            image=img,
-            model=model,
-            result=result,
-            output_file=output_file,
-            window_name=backend.value,
-            show_result=show_result)
+        if not isinstance(img, Sequence):
+            img = [img]
+        for single_img in img:
+            task_processor.visualize(
+                image=single_img,
+                model=model,
+                result=result,
+                output_file=output_file,
+                window_name=backend.value,
+                show_result=show_result)
     except Exception as e:
         from mmdeploy.utils import get_root_logger
         logger = get_root_logger()

mmdeploy/codebase/mmdet/models/roi_heads/test_mixins.py

@@ -96,7 +96,7 @@ def mask_test_mixin__simple_test_mask(ctx, self, x, img_metas, det_bboxes,
     # expand might lead to static shape, use broadcast instead
     batch_index = torch.arange(
         det_bboxes.size(0), device=det_bboxes.device).float().view(
-            -1, 1) + det_bboxes.new_zeros(
+            -1, 1, 1) + det_bboxes.new_zeros(
                 (det_bboxes.size(0), det_bboxes.size(1))).unsqueeze(-1)
     mask_rois = torch.cat([batch_index, det_bboxes], dim=-1)
     mask_rois = mask_rois.view(-1, 5)

tools/deploy.py

@@ -27,7 +27,11 @@ def parse_args():
     parser.add_argument('checkpoint', help='model checkpoint path')
     parser.add_argument('img', help='image used to convert model model')
     parser.add_argument(
-        '--test-img', default=None, help='image used to test model')
+        '--test-img',
+        default=None,
+        type=str,
+        nargs='+',
+        help='image used to test model')
     parser.add_argument(
         '--work-dir',
         default=os.getcwd(),