From 13c787261919cff9e05b18bc0a22d8005a518853 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 15:05:06 -0400 Subject: [PATCH 01/10] Add torchvision.utils.draw_bounding_boxes This is entirely untested. --- torchvision/utils.py | 61 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/torchvision/utils.py b/torchvision/utils.py index be373138c5f..c32e07adbc4 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -128,3 +128,64 @@ def save_image( ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy() im = Image.fromarray(ndarr) im.save(fp, format=format) + + +BBox = Tuple[int, int, int, int] +BBoxes = Sequence[BBox] +Color = Tuple[int, int, int] +DEFAULT_COLORS: Sequence[Color] + + +def draw_bounding_boxes( + image: torch.Tensor, + bboxes: Union[BBoxes, Dict[str, Sequence[BBox]]], + colors: Optional[Dict[str, Color]] = None, + draw_labels: bool = None, + width: int = 1, +) -> torch.Tensor: + # TODO: docstring + + bboxes_is_seq = BBoxes.__instancecheck__(bboxes) + # bboxes_is_dict is Dict[str, Sequence[BBox]].__instancecheck__(bboxes) + bboxes_is_dict = not bboxes_is_seq + + if bboxes_is_seq: + # TODO: raise better Errors + if colors is not None: + # can't pass custom colors if bboxes is a sequence + raise Error + if draw_labels is True: + # can't draw labels if bboxes is a sequence + raise Error + + if draw_labels is None: + if bboxes_is_seq: + draw_labels = False + else: # BBoxes.__instancecheck__(Dict[str, Sequence[BBox]]) + draw_labels = True + + if colors is None: + # TODO: default to one of @pmeir's suggestions + pass + + from PIL import Image, ImageDraw + # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer + ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute( + 1, 2, 0).to('cpu', torch.uint8).numpy() + im = Image.fromarray(ndarr) + draw = ImageDraw.Draw(im) + + if bboxes_is_dict: + for bbox_class, bbox in bboxes.items(): + if draw_labels: + # TODO: this will probably overlap with the bbox + # hard-code in a margin for the label? + label_tl_x, label_tl_y, _, _ = bbox + draw.text((label_tl_x, label_tl_y), bbox_class) + draw.rectangle(bbox, outline=color, width=width) + else: + for bbox, color in zip(bboxes, colors): + draw.rectangle(bbox, outline=color, width=width) + + from numpy import array as to_numpy_array + return torch.from_numpy(to_numpy_array(im)) From 615f836ff31525eb5fa28d325e4dc693bd242756 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 15:15:34 -0400 Subject: [PATCH 02/10] Add test_bboxes_not_inplace This is entirely untested. --- test/test_utils.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/test_utils.py b/test/test_utils.py index f1982130f75..a82c47ff400 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -40,6 +40,23 @@ def test_normalize_in_make_grid(self): self.assertTrue(torch.equal(norm_max, rounded_grid_max), 'Normalized max is not equal to 1') self.assertTrue(torch.equal(norm_min, rounded_grid_min), 'Normalized min is not equal to 0') + + def test_bboxes_not_inplace(self): + t = torch.rand(5, 3, 10, 10) * 255 + t_clone = t.clone() + + utils.draw_bounding_bboxes(t, draw_labels=False) + self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') + + utils.draw_bounding_bboxes(t, draw_labels=True) + self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') + + def test_bboxes(self): + IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") + IMAGE_DIR = os.path.join(IMAGE_ROOT, "fakedata", "imagefolder") + img_path = os.path.join(IMAGE_DIR, 'a4.png') + img_pil = torch.from_numpy(np.array(Image.open(img_path))) + # TODO: implement this test @unittest.skipIf('win' in sys.platform, 'temporarily disabled on Windows') def test_save_image(self): From da4773faf9aa6c27a4bd8da39cfd4283786c52a1 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 15:21:24 -0400 Subject: [PATCH 03/10] Import Sequence type in utils.draw_bounding_boxes --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index c32e07adbc4..3d31ba8659b 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -1,4 +1,4 @@ -from typing import Union, Optional, List, Tuple, Text, BinaryIO +from typing import Union, Optional, List, Tuple, Text, BinaryIO, Sequence import io import pathlib import torch From aba0b47ba6aa27f3288836cc2e979c021edc42fd Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 15:34:27 -0400 Subject: [PATCH 04/10] Import Dict type in utils.draw_bounding_boxes --- torchvision/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 3d31ba8659b..b2b73df03bf 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -1,4 +1,4 @@ -from typing import Union, Optional, List, Tuple, Text, BinaryIO, Sequence +from typing import Union, Optional, List, Tuple, Text, BinaryIO, Sequence, Dict import io import pathlib import torch From 7d97ebc4abf295650ea8d4c8d71924ae4ec79b39 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 15:41:51 -0400 Subject: [PATCH 05/10] Add test_bboxes --- test/test_utils.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index a82c47ff400..04ad436cb0a 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -40,23 +40,38 @@ def test_normalize_in_make_grid(self): self.assertTrue(torch.equal(norm_max, rounded_grid_max), 'Normalized max is not equal to 1') self.assertTrue(torch.equal(norm_min, rounded_grid_min), 'Normalized min is not equal to 0') - - def test_bboxes_not_inplace(self): - t = torch.rand(5, 3, 10, 10) * 255 - t_clone = t.clone() - utils.draw_bounding_bboxes(t, draw_labels=False) - self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') + # def test_bboxes_not_inplace(self): + # t = torch.rand(5, 3, 10, 10) * 255 + # t_clone = t.clone() + # + # TODO: this doesn't work; we need to pass in bboxes + # utils.draw_bounding_bboxes(t, draw_labels=False) + # self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') + # + # utils.draw_bounding_bboxes(t, draw_labels=True) + # self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') - utils.draw_bounding_bboxes(t, draw_labels=True) - self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') - def test_bboxes(self): IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") IMAGE_DIR = os.path.join(IMAGE_ROOT, "fakedata", "imagefolder") - img_path = os.path.join(IMAGE_DIR, 'a4.png') - img_pil = torch.from_numpy(np.array(Image.open(img_path))) - # TODO: implement this test + inp_img_path = os.path.join(IMAGE_DIR, 'a4.png') + out_img_path = os.path.join(IMAGE_DIR, 'b5.png') + + inp_img_pil = torch.from_numpy(np.array(Image.open(inp_img_path))) + bboxes = ((1, 2, 10, 18), (4, 8, 9, 11)) + # TODO: maybe write the rectangle programatically in this test instead of + # statically loading output? + out_img_pil = torch.from_numpy(np.array(Image.open(out_img_path))) + + self.assertTrue( + torch.equal( + utils.draw_bounding_bboxes(inp_img_pil, bboxes, draw_labels=False), + out_img_pil, + ), + 'draw_bounding_bboxes returned an incorrect result', + ) + @unittest.skipIf('win' in sys.platform, 'temporarily disabled on Windows') def test_save_image(self): From 530bc59addbb8504a5e15717e150e4e23ec6bd6b Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 16:26:19 -0400 Subject: [PATCH 06/10] Fix colors kwarg in utils.draw_bboxes --- torchvision/utils.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index b2b73df03bf..d09a798ea5b 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -165,7 +165,7 @@ def draw_bounding_boxes( draw_labels = True if colors is None: - # TODO: default to one of @pmeir's suggestions + # TODO: default to one of @pmeir's suggestions as a seq pass from PIL import Image, ImageDraw @@ -176,16 +176,20 @@ def draw_bounding_boxes( draw = ImageDraw.Draw(im) if bboxes_is_dict: - for bbox_class, bbox in bboxes.items(): + if Sequence[Color].__instancecheck__(colors): + # align the colors seq with the bbox classes + colors = dict(zip(sorted(bboxes.keys()), colors)) + + for i, (bbox_class, bbox) in enumerate(bboxes.items()): + draw.rectangle(bbox, outline=colors[bbox_class], width=width) if draw_labels: # TODO: this will probably overlap with the bbox # hard-code in a margin for the label? label_tl_x, label_tl_y, _, _ = bbox draw.text((label_tl_x, label_tl_y), bbox_class) - draw.rectangle(bbox, outline=color, width=width) - else: - for bbox, color in zip(bboxes, colors): - draw.rectangle(bbox, outline=color, width=width) + else: # bboxes_is_seq + for i, bbox in enumerate(bboxes): + draw.rectangle(bbox, outline=colors[i], width=width) from numpy import array as to_numpy_array return torch.from_numpy(to_numpy_array(im)) From 4afb3d7f89ee84d58d9f1f447954018d13edd4ce Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 16:26:55 -0400 Subject: [PATCH 07/10] Import np.array in test_utils.test_bboxes --- test/test_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index 04ad436cb0a..4855d2daff9 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -53,6 +53,8 @@ def test_normalize_in_make_grid(self): # self.assertTrue(torch.equal(t, t_clone), 'draw_bounding_bboxes modified tensor in-place') def test_bboxes(self): + from numpy import array as to_numpy_array + IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") IMAGE_DIR = os.path.join(IMAGE_ROOT, "fakedata", "imagefolder") inp_img_path = os.path.join(IMAGE_DIR, 'a4.png') @@ -62,7 +64,7 @@ def test_bboxes(self): bboxes = ((1, 2, 10, 18), (4, 8, 9, 11)) # TODO: maybe write the rectangle programatically in this test instead of # statically loading output? - out_img_pil = torch.from_numpy(np.array(Image.open(out_img_path))) + out_img_pil = torch.from_numpy(to_numpy_array(Image.open(out_img_path))) self.assertTrue( torch.equal( @@ -71,7 +73,6 @@ def test_bboxes(self): ), 'draw_bounding_bboxes returned an incorrect result', ) - @unittest.skipIf('win' in sys.platform, 'temporarily disabled on Windows') def test_save_image(self): From e8b2f8c06265cc74386f4d055b487a7097f6f694 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 16:36:05 -0400 Subject: [PATCH 08/10] Rename tensor to image in draw_bounding_boxes --- torchvision/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index d09a798ea5b..64d61c907a2 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -163,14 +163,14 @@ def draw_bounding_boxes( draw_labels = False else: # BBoxes.__instancecheck__(Dict[str, Sequence[BBox]]) draw_labels = True - + if colors is None: # TODO: default to one of @pmeir's suggestions as a seq pass - + from PIL import Image, ImageDraw # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer - ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute( + ndarr = image.mul(255).add_(0.5).clamp_(0, 255).permute( 1, 2, 0).to('cpu', torch.uint8).numpy() im = Image.fromarray(ndarr) draw = ImageDraw.Draw(im) @@ -179,12 +179,12 @@ def draw_bounding_boxes( if Sequence[Color].__instancecheck__(colors): # align the colors seq with the bbox classes colors = dict(zip(sorted(bboxes.keys()), colors)) - + for i, (bbox_class, bbox) in enumerate(bboxes.items()): draw.rectangle(bbox, outline=colors[bbox_class], width=width) if draw_labels: # TODO: this will probably overlap with the bbox - # hard-code in a margin for the label? + # hard-code in a margin for the label? label_tl_x, label_tl_y, _, _ = bbox draw.text((label_tl_x, label_tl_y), bbox_class) else: # bboxes_is_seq From a90255766a468d93c3466e57b839bf9948181e49 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 16:54:34 -0400 Subject: [PATCH 09/10] Import np.array in test_utils.test_bboxes v2 --- test/test_utils.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index 4855d2daff9..f1a8d65db6d 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -9,6 +9,11 @@ from PIL import Image +def path_to_tensor(filepath): + from numpy import array as to_numpy_array + return torch.from_numpy(to_numpy_array(Image.open(filepath))) + + class Tester(unittest.TestCase): def test_make_grid_not_inplace(self): @@ -60,11 +65,11 @@ def test_bboxes(self): inp_img_path = os.path.join(IMAGE_DIR, 'a4.png') out_img_path = os.path.join(IMAGE_DIR, 'b5.png') - inp_img_pil = torch.from_numpy(np.array(Image.open(inp_img_path))) + inp_img_pil = path_to_tensor(inp_img_path) bboxes = ((1, 2, 10, 18), (4, 8, 9, 11)) # TODO: maybe write the rectangle programatically in this test instead of # statically loading output? - out_img_pil = torch.from_numpy(to_numpy_array(Image.open(out_img_path))) + out_img_pil = path_to_tensor(out_img_path) self.assertTrue( torch.equal( From 2ec6a2dc7063d4d8d1fc536e2976857483006ec9 Mon Sep 17 00:00:00 2001 From: Sumanth Ratna Date: Sat, 29 Aug 2020 18:31:36 -0400 Subject: [PATCH 10/10] Fix colors type --- torchvision/utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 64d61c907a2..54c5278aef0 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -164,9 +164,12 @@ def draw_bounding_boxes( else: # BBoxes.__instancecheck__(Dict[str, Sequence[BBox]]) draw_labels = True + # colors: Union[Sequence[Color], Dict[str, Color]] if colors is None: # TODO: default to one of @pmeir's suggestions as a seq - pass + colors_: Sequence[Color] = colors + else: + colors_: Dict[str, Color] = colors from PIL import Image, ImageDraw # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer @@ -176,12 +179,12 @@ def draw_bounding_boxes( draw = ImageDraw.Draw(im) if bboxes_is_dict: - if Sequence[Color].__instancecheck__(colors): + if Sequence[Color].__instancecheck__(colors_): # align the colors seq with the bbox classes - colors = dict(zip(sorted(bboxes.keys()), colors)) + colors = dict(zip(sorted(bboxes.keys()), colors_)) - for i, (bbox_class, bbox) in enumerate(bboxes.items()): - draw.rectangle(bbox, outline=colors[bbox_class], width=width) + for bbox_class, bbox in enumerate(bboxes.items()): + draw.rectangle(bbox, outline=colors_[bbox_class], width=width) if draw_labels: # TODO: this will probably overlap with the bbox # hard-code in a margin for the label? @@ -189,7 +192,7 @@ def draw_bounding_boxes( draw.text((label_tl_x, label_tl_y), bbox_class) else: # bboxes_is_seq for i, bbox in enumerate(bboxes): - draw.rectangle(bbox, outline=colors[i], width=width) + draw.rectangle(bbox, outline=colors_[i], width=width) from numpy import array as to_numpy_array return torch.from_numpy(to_numpy_array(im))