fixes mmdet tensor_to_image (#804)

airctic · Apr 19, 2021 · f2257cd · f2257cd
1 parent 0d0bdb4
commit f2257cd
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 9 deletions.
diff --git a/icevision/models/mmdet/common/bbox/prediction.py b/icevision/models/mmdet/common/bbox/prediction.py
@@ -10,6 +10,7 @@
 from icevision.core import *
 from icevision.data import *
 from icevision.models.utils import _predict_from_dl
+from icevision.models.mmdet.common.utils import *
 from icevision.models.mmdet.common.bbox.dataloaders import build_infer_batch
 from icevision.models.mmdet.common.utils import convert_background_from_last_to_zero
 
@@ -137,8 +138,7 @@ def convert_raw_prediction(
     pred.above_threshold = keep_mask
 
     if keep_image:
-        image = sample["img"]
-        image = image.detach().cpu().numpy().transpose(1, 2, 0)
+        image = mmdet_tensor_to_image(sample["img"])
 
         pred.set_img(image)
         record.set_img(image)

diff --git a/icevision/models/mmdet/common/bbox/show_batch.py b/icevision/models/mmdet/common/bbox/show_batch.py
@@ -2,6 +2,7 @@
 
 from icevision.utils import *
 from icevision.visualize import *
+from icevision.models.mmdet.common.utils import *
 
 
 def show_batch(batch_and_records, ncols: int = 1, figsize=None, **show_samples_kwargs):
@@ -20,9 +21,7 @@ def show_batch(batch_and_records, ncols: int = 1, figsize=None, **show_samples_k
         }
 
     for tensor_image, record in zip(batch["img"][:], records):
-        image = (
-            tensor_image.detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1].copy()
-        )
+        image = mmdet_tensor_to_image(tensor_image)
         record.set_img(image)
 
     return show_samples(records, ncols=ncols, figsize=figsize, **show_samples_kwargs)
diff --git a/icevision/models/mmdet/common/mask/prediction.py b/icevision/models/mmdet/common/mask/prediction.py
@@ -10,7 +10,7 @@
 from icevision.core import *
 from icevision.data import *
 from icevision.models.utils import _predict_from_dl
-from icevision.models.mmdet.common.utils import convert_background_from_last_to_zero
+from icevision.models.mmdet.common.utils import *
 from icevision.models.mmdet.common.mask.dataloaders import *
 from icevision.models.mmdet.common.bbox.prediction import (
     _unpack_raw_bboxes,
@@ -142,8 +142,7 @@ def convert_raw_prediction(
     pred.above_threshold = keep_mask
 
     if keep_image:
-        image = sample["img"]
-        image = image.detach().cpu().numpy().transpose(1, 2, 0)
+        image = mmdet_tensor_to_image(sample["img"])
 
         pred.set_img(image)
         record.set_img(image)

diff --git a/icevision/models/mmdet/common/mask/show_batch.py b/icevision/models/mmdet/common/mask/show_batch.py
@@ -2,6 +2,7 @@
 
 from icevision.utils import *
 from icevision.visualize import *
+from icevision.models.mmdet.common.utils import *
 
 
 def show_batch(batch_and_records, ncols: int = 1, figsize=None, **show_samples_kwargs):
@@ -12,7 +13,7 @@ def show_batch(batch_and_records, ncols: int = 1, figsize=None, **show_samples_k
     batch, records = batch_and_records
 
     for tensor_image, record in zip(batch["img"][:], records):
-        image = tensor_image.detach().cpu().numpy().transpose(1, 2, 0)
+        image = mmdet_tensor_to_image(tensor_image)
         record.set_img(image)
 
     return show_samples(records, ncols=ncols, figsize=figsize, **show_samples_kwargs)
diff --git a/icevision/models/mmdet/common/utils.py b/icevision/models/mmdet/common/utils.py
@@ -1,10 +1,12 @@
 __all__ = [
     "convert_background_from_zero_to_last",
     "convert_background_from_last_to_zero",
+    "mmdet_tensor_to_image",
     "build_model",
 ]
 
 from icevision.imports import *
+from icevision.utils import *
 from mmcv import Config
 from mmdet.models import build_detector
 from mmcv.runner import load_checkpoint
@@ -23,6 +25,10 @@ def convert_background_from_last_to_zero(label_ids, class_map):
     return label_ids
 
 
+def mmdet_tensor_to_image(tensor_image):
+    return tensor_to_image(tensor_image)[:, :, ::-1].copy()
+
+
 def build_model(
     model_type: str,
     backbone: MMDetBackboneConfig,