[Fix] fix empty image exporting bug (open-mmlab#196)

* fix empty image exporting bug

* solve black empty image

* concat zeros ahead

* add docstring

mmdeploy/mmdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py

@@ -94,10 +94,15 @@ def single_roi_extractor__forward(ctx,
                                   feats,
                                   rois,
                                   roi_scale_factor=None):
-    """Rewrite `forward` for default backend.
+    """Rewrite `forward` of SingleRoIExtractor for default backend.
 
-    Add mark for roi_extractor forward. Remove unnecessary code of origin
-    forward function.
+    Rewrite this function to enable exporting to onnx even though the input
+    image contains no targets. Note that, `ScatterND` of onnx may conflict with
+    `Reshape` if a tensor have a dim size of 0. Thus, we have to cat zeros to
+    the dim 0 of `roi_feats` and recover back after all roi align finished.
+
+    Besides, this function adds mark for roi_extractor forward and remove
+    unnecessary code of origin forward function.
     """
     out_size = self.roi_layers[0].output_size
     num_levels = len(feats)
@@ -111,17 +116,29 @@ def single_roi_extractor__forward(ctx,
     if roi_scale_factor is not None:
         rois = self.roi_rescale(rois, roi_scale_factor)
 
+    # concat len num_levels * 2 of zero tensors to dim 0 of roi_feats
+    roi_feats = torch.cat(
+        (roi_feats.new_zeros(num_levels * 2,
+                             *roi_feats.shape[-3:]), roi_feats))
     for i in range(num_levels):
         mask = target_lvls == i
         inds = mask.nonzero(as_tuple=False).squeeze(1)
 
-        # expand tensor to eliminate [0, ...] tensor
-        rois_i = torch.cat((rois[inds], rois.new_zeros(1, 5)))
+        # concat len 2 zero tensors to dim 0 of roi_feats
+        rois_i = torch.cat((rois.new_zeros(2, 5), rois[inds]))
 
         roi_feats_t = self.roi_layers[i](feats[i], rois_i)
 
-        # slice and recover the tensor
-        roi_feats[inds] = roi_feats_t[0:-1]
+        # correspondingly change the inds
+        inds = torch.cat([
+            torch.tensor([2 * i, 2 * i + 1],
+                         device=inds.device,
+                         dtype=inds.dtype), inds + num_levels * 2
+        ])
+        roi_feats[inds] = roi_feats_t
+
+    # slice and recover tensors
+    roi_feats = roi_feats[num_levels * (2):]
     return roi_feats
 
 

mmdeploy/mmdet/models/roi_heads/standard_roi_head.py

@@ -1,3 +1,5 @@
+import torch
+
 from mmdeploy.core import FUNCTION_REWRITER
 
 
@@ -37,6 +39,21 @@ def standard_roi_head__simple_test(ctx, self, x, proposals, img_metas,
     if not self.with_mask:
         return det_bboxes, det_labels
 
+    # padding zeros to det_bboxes and det_labels
+    det_bboxes_tail = torch.zeros(
+        det_bboxes.size(0),
+        1,
+        det_bboxes.size(2),
+        device=det_bboxes.device,
+        dtype=det_bboxes.dtype)
+    det_labels_tail = torch.zeros(
+        det_labels.size(0),
+        1,
+        device=det_labels.device,
+        dtype=det_labels.dtype)
+    det_bboxes = torch.cat([det_bboxes, det_bboxes_tail], 1)
+    det_labels = torch.cat([det_labels, det_labels_tail], 1)
+
     segm_results = self.simple_test_mask(
         x, img_metas, det_bboxes, det_labels, rescale=False)
     return det_bboxes, det_labels, segm_results

mmdeploy/mmdet/models/roi_heads/test_mixins.py

@@ -82,16 +82,6 @@ def mask_test_mixin__simple_test_mask(ctx, self, x, img_metas, det_bboxes,
         tuple[Tensor]: (segm_results), `segm_results` of shape
         [N, num_det, roi_H, roi_W].
     """
-    if det_bboxes.shape[1] == 0:
-        bboxes_shape, labels_shape = list(det_bboxes.shape), list(
-            det_labels.shape)
-        bboxes_shape[1], labels_shape[1] = 1, 1
-        det_bboxes = torch.tensor(
-            [[[0., 0., 1., 1., 0.]]],
-            device=det_bboxes.device).expand(*bboxes_shape)
-        det_labels = torch.tensor(
-            [[0]], device=det_bboxes.device).expand(*labels_shape)
-
     batch_size = det_bboxes.size(0)
     det_bboxes = det_bboxes[..., :4]
     batch_index = torch.arange(