From 97226fe8f3f495821a35e112dbe6945bab9248d1 Mon Sep 17 00:00:00 2001 From: Sugato Ray Date: Thu, 10 Jun 2021 13:03:43 -0500 Subject: [PATCH 1/2] allow single-channel-images in draw_bounding_boxes This is bug-fix for issue #4042. https://github.com/pytorch/vision/issues/4042 --- torchvision/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torchvision/utils.py b/torchvision/utils.py index ea197d7386e..5525477561e 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -180,6 +180,10 @@ def draw_bounding_boxes( raise ValueError("Pass individual images, not batches") ndarr = image.permute(1, 2, 0).numpy() + # allow single-channel-images + # shape: (1, H, W) with C = 1 + if ndarr.shape[-1]==1: + ndarr = np.tile(ndarr, (1, 1, 3)) img_to_draw = Image.fromarray(ndarr) img_boxes = boxes.to(torch.int64).tolist() From 8a384cb9a5eace3f90aa01fe11e58867b5aa5dba Mon Sep 17 00:00:00 2001 From: Sugato Ray Date: Thu, 10 Jun 2021 14:06:06 -0500 Subject: [PATCH 2/2] change made for passing linting requirements #4042 --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 5525477561e..90657acee4f 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -182,7 +182,7 @@ def draw_bounding_boxes( ndarr = image.permute(1, 2, 0).numpy() # allow single-channel-images # shape: (1, H, W) with C = 1 - if ndarr.shape[-1]==1: + if ndarr.shape[-1] == 1: ndarr = np.tile(ndarr, (1, 1, 3)) img_to_draw = Image.fromarray(ndarr)