Fix: Improve the bounding boxes implementation (#3075)

* Fix: Improve the bounding boxes implementation

Use write_png instead of PIL in test_draw_boxes()
Initialize txt_font only once

* Remove channels permutation in test_draw_boxes
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>
Co-authored-by: Aditya Khursale <akhursale@nvidia.com>
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>

test/test_utils.py

@@ -6,7 +6,7 @@ import torchvision.utils as utils
 import unittest
 from io import BytesIO
 import torchvision.transforms.functional as F
-from torchvision.io.image import read_image
+from torchvision.io.image import read_image, write_png
 from PIL import Image
 
 
@@ -90,7 +90,7 @@ class Tester(unittest.TestCase):
 
         path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_util.png")
         if not os.path.exists(path):
-            Image.fromarray(result.permute(1, 2, 0).numpy()).save(path)
+            write_png(result, path)
 
         expected = read_image(path)
         self.assertTrue(torch.equal(result, expected))

torchvision/utils.py

@@ -177,13 +177,13 @@ def draw_bounding_boxes(
     img_boxes = boxes.to(torch.int64).tolist()
 
     draw = ImageDraw.Draw(img_to_draw)
+    txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size)
 
     for i, bbox in enumerate(img_boxes):
         color = None if colors is None else colors[i]
         draw.rectangle(bbox, width=width, outline=color)
 
         if labels is not None:
-            txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size)
             draw.text((bbox[0], bbox[1]), labels[i], fill=color, font=txt_font)
 
     return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1)