Fix no object for segment

Handle case when no bounding boxes are detected during segmentation

infer-seg.py

@@ -42,6 +42,11 @@ def main(args: argparse.Namespace) -> None:
                                 device=device)
         bboxes, scores, labels, masks = seg_postprocess(
             data, bgr.shape[:2], args.conf_thres, args.iou_thres)
+        if bboxes is None:
+            # if no bounding box or others save original image
+            if not args.show:
+                cv2.imwrite(str(save_image), draw)
+            continue
         masks = masks[:, dh:H - dh, dw:W - dw, :]
         indices = (labels % len(MASK_COLORS)).long()
         mask_colors = torch.asarray(MASK_COLORS, device=device)[indices]

models/torch_utils.py

@@ -18,6 +18,8 @@ def seg_postprocess(
     bboxes, scores, labels, maskconf = outputs.split([4, 1, 1, 32], 1)
     scores, labels = scores.squeeze(), labels.squeeze()
     idx = scores > conf_thres
+    if idx.sum() == 0:  # no bounding boxes or seg were created
+        return None, None, None, None
     bboxes, scores, labels, maskconf = \
         bboxes[idx], scores[idx], labels[idx], maskconf[idx]
     idx = batched_nms(bboxes, scores, labels, iou_thres)