bounding boxes in torchvision.utils

[sumanthratna]

🚀 Feature

I'd like to easily be able to draw a bounding box onto an image represented as a torch tensor.

Motivation

Using YOLO, I get a bunch of bounding boxes in an image. I want to be able to easily draw those onto a torch tensor (created via torchvision.transforms.ToTensor). This seems like a reasonable request because other bbox utilities such as NMS exist in torchvision.

Pitch

tensor_image = torch.tensor(...)
new_img = torchvision.utils.draw_bounding_box(
    tensor_image,
    [x_min, y_min, x_max, y_max],
)

Alternatives

I think the following works right now, but the torch.Tensor > PIL.Image > torch.Tensor conversion for a single operation per image makes me feel uncomfortable due to efficiency.

# entirely untested!

tensor_image = torch.tensor(...)
def draw_bounding_box(tensor, bbox, **kwargs):  # TODO: don't use a kwargs dict
    # adapted from https://pytorch.org/docs/stable/_modules/torchvision/utils.html#save_image
    from PIL import Image
    # TODO: none of the kwargs to make_grid are in this scope
    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                     normalize=normalize, range=range, scale_each=scale_each)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()

    # TODO: convert bbox to the appropriate type
    return cv2.rectangle(ndarr, bbox)  # use PIL if opencv-python isn't a dependency

Additional context

See cv2.rectangle (source) and PIL.ImageDraw.rectangle (source).

[pmeier]

Hey @sumanthratna this sounds like a reasonable request. Lets wait for @fmassa to approve this, but I think we can move forward with this. Would you like to send a PR?

[sumanthratna]

Would you like to send a PR?

Sure @pmeier!

Here are the solutions I'm thinking of:

• converting the torch tensor to a numpy array or PIL Image (see "Alternatives")
  • pro: This will make drawing the actual bbox a job for either cv2 or PIL as opposed to torchvision, which means less maintenance work for you
  • con: if I use a for-loop to call this new torchvision function multiple times on the same image, converting to a np array or PIL Image for each iteration induces unnecessary torch.Tensor > np.array ( > PIL.Image) overhead
• porting cv2.rectangle to torch
  • pro: for-loops won't make this function particularly inefficient

What do you think is the best option? I'm personally leaning towards the second option.

[fmassa]

Hi,

We do want to have functions for drawing bounding boxes in images, but we need to come up with requirements and API. If we plot a rectangle, we would probably also want to support plotting text, selecting colors, line width, etc. Additionally, should the API support a single box at a time or a batch of boxes?
All those questions I'm not yet clear about, but would be good to be thought about and discussed before we start implementing anything.

torchvision currently don't depend on OpenCV, and we would like to keep it that way. So the box drawing function would need to depend on either: PIL, matplotlib or be natively implemented in torchvision

I would not be very concerned about the overhead of this function (converting a Tensor to an ndarray is cheap, and the same for a PIL.Image), plus I wouldn't expect this to be on a critical path so it's ok if it's not super efficient.

[pmeier]

We do want to have functions for drawing bounding boxes in images, but we need to come up with requirements and API

If we only go for drawing boxes, I think this can be done straight forward in torchvision. We basically only need to calculate slices for the current box and just set the pixels to the given color.

we would probably also want to support plotting text

This will be a problem in torchvision. If we also want this I suggest we forget doing this natively and use ImageDraw from PIL.

Additionally, should the API support a single box at a time or a batch of boxes?

I'm voting for multiple boxes to avoid a color conflict between boxes. If we only plot one box at a time either the user has to manually specify a new color for each box or we need to somehow track which colors we have already used.

[sumanthratna]

@fmassa those are good questions that I didn't consider, and I definitely agree with what @pmeier has suggested.

Additionally, should the API support a single box at a time or a batch of boxes?

I'm voting for multiple boxes to avoid a color conflict between boxes. If we only plot one box at a time either the user has to manually specify a new color for each box or we need to somehow track which colors we have already used.

Originally, I thought that we should leave it up to users to decide the color, to avoid any confusing behavior. However, I do think that batch drawing might make for cleaner code on the user-end. All in all, I think we should go with batch drawing. I'm still unsure if we want to deal with color conflicts, though. Sometimes a user will want to draw multiple boxes with the same class (see the image below, source):

Here's what I've written as a workaround until this gets merged (I haven't tested it):

def draw_bounding_box(
    tensor,
    bbox,
    nrow=8,
    padding=2,
    normalize=False,
    range=None,
    scale_each=False,
    pad_value=0,
    fill=None,
    outline=None,
    width=1,
):
    from torchvision.utils import make_grid
    from PIL import Image, ImageDraw
    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                     normalize=normalize, range=range, scale_each=scale_each)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(
        1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)

    ImageDraw.Draw(im).rectangle(bbox, fill=fill, outline=outline, width=width)
    from numpy import array as to_numpy_array
    return torch.from_numpy(to_numpy_array(im))

It should be simple to convert this to a batch-bbox function.

[pmeier]

@sumanthratna I don't think we should couple this with make_grid. Can't we just have draw_bounding_box and if you need a grid of images pass the output to make_grid afterwards?

[sumanthratna]

That seems reasonable. I initially just copied torchvision.utils.save_image without paying attention to what it does :)

Is this better? We might not need to unnormalize.

def draw_bounding_box(
    tensor,
    bbox,
    fill=None,
    outline=None,
    width=1,
):
    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)

    ImageDraw.Draw(im).rectangle(bbox, fill=fill, outline=outline, width=width)
    from numpy import array as to_numpy_array
    return torch.from_numpy(to_numpy_array(im))

[pmeier]

1. I don't think we should have the fill parameter.
2. I would rename outline to color to make it more expressive.
3. Right now the function is only taking a single bounding box. We need to extend this, but I'm not sure how make this intuitively. Our solution should not only support multiple boxes, but also multiple boxes of the same class.

[sumanthratna]

1. Why not? If PIL.ImageDraw.rectangle accepts the fill kwarg, why would we not want to use it? IMHO we don't need to limit the functionality of Pillow
2. This makes sense if we remove the fill kwarg.
3. Definitely. It seems that if we want to deal with classes, a dict of {detection_class: [bbox1, bbox2, ...]} seems the most intuitive way to deal with this. Alternatively, we could just make our bbox arg a sequence of bboxes, and tell users to call draw_bounding_box multiple times (1 call for each class).

[pmeier]

Why not?

As the name implies, fill will fill the inside of the rectangle and thus blocking everything within the bounding box from view. I can't think of a scenario where someone wants that.