def test_draw_no_segmention_mask():
img = torch.full((3, 100, 100), 0, dtype=torch.uint8)
masks = torch.full((0, 100, 100), 0, dtype=torch.bool)
with pytest.warns(
UserWarning,
match=re.escape(
"masks doesn't contain any mask. No mask was drawn")):
res = utils.draw_segmentation_masks(img, masks)
# Check that the function didnt change the image
assert res.eq(img).all()
def test_draw_segmentation_masks_no_colors(self):
img = torch.full((3, 20, 20), 255, dtype=torch.uint8)
result = utils.draw_segmentation_masks(img, masks, colors=None)
path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"assets", "fakedata",
"draw_segm_masks_no_colors_util.png")
if not os.path.exists(path):
res = Image.fromarray(result.permute(1, 2, 0).contiguous().numpy())
res.save(path)
expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
self.assertTrue(torch.equal(result, expected))
def test_draw_segmentation_masks(colors, alpha):
"""This test makes sure that masks draw their corresponding color where they should"""
num_masks, h, w = 2, 100, 100
dtype = torch.uint8
img = torch.randint(0, 256, size=(3, h, w), dtype=dtype)
masks = torch.randint(0, 2, (num_masks, h, w), dtype=torch.bool)
# For testing we enforce that there's no overlap between the masks. The
# current behaviour is that the last mask's color will take priority when
# masks overlap, but this makes testing slightly harder so we don't really
# care
overlap = masks[0] & masks[1]
masks[:, overlap] = False
out = utils.draw_segmentation_masks(img, masks, colors=colors, alpha=alpha)
assert out.dtype == dtype
assert out is not img
# Make sure the image didn't change where there's no mask
masked_pixels = masks[0] | masks[1]
assert_equal(img[:, ~masked_pixels], out[:, ~masked_pixels])
if colors is None:
colors = utils._generate_color_palette(num_masks)
elif isinstance(colors, str) or isinstance(colors, tuple):
colors = [colors]
# Make sure each mask draws with its own color
for mask, color in zip(masks, colors):
if isinstance(color, str):
color = ImageColor.getrgb(color)
color = torch.tensor(color, dtype=dtype)
if alpha == 1:
assert (out[:, mask] == color[:, None]).all()
elif alpha == 0:
assert (out[:, mask] == img[:, mask]).all()
interpolated_color = (img[:, mask] * (1 - alpha) +
color[:, None] * alpha).to(dtype)
torch.testing.assert_close(out[:, mask],
interpolated_color,
rtol=0.0,
atol=1.0)
def test_draw_segmentation_masks_colors(self):
img = torch.full((3, 5, 5), 255, dtype=torch.uint8)
img_cp = img.clone()
masks_cp = masks.clone()
colors = ["#FF00FF", (0, 255, 0), "red"]
result = utils.draw_segmentation_masks(img, masks, colors=colors)
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
"fakedata", "draw_segm_masks_colors_util.png")
if not os.path.exists(path):
res = Image.fromarray(result.permute(1, 2, 0).contiguous().numpy())
res.save(path)
expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
self.assertTrue(torch.equal(result, expected))
# Check if modification is not in place
self.assertTrue(torch.all(torch.eq(img, img_cp)).item())
self.assertTrue(torch.all(torch.eq(masks, masks_cp)).item())
def visualize_batch(self):
loader = DataLoader(self, shuffle=True, batch_size=4)
if self.mode != "test":
imgs, masks, fnames = next(iter(loader))
else:
imgs, fnames = next(iter(loader))
batch_inputs = F.convert_image_dtype(imgs, dtype=torch.uint8)
if self.mode != "test":
batch_outputs = F.convert_image_dtype(masks, dtype=torch.bool)
list_imgs = [
draw_segmentation_masks(img,
masks=mask,
alpha=0.6,
colors=(102, 255, 178))
for img, mask in zip(batch_inputs, batch_outputs)
]
else:
list_imgs = [imgs[i] for i in range(len(imgs))]
self.show(list_imgs, fnames)
# can read it as the following query: "For which pixels is 'dog' the most likely
# class?"
#
# .. note::
# While we're using the ``normalized_masks`` here, we would have
# gotten the same result by using the non-normalized scores of the model
# directly (as the softmax operation preserves the order).
#
# Now that we have boolean masks, we can use them with
# :func:`~torchvision.utils.draw_segmentation_masks` to plot them on top of the
# original images:
from torchvision.utils import draw_segmentation_masks
dogs_with_masks = [
draw_segmentation_masks(img, masks=mask, alpha=0.7)
for img, mask in zip(batch_int, boolean_dog_masks)
]
show(dogs_with_masks)
#####################################
# We can plot more than one mask per image! Remember that the model returned as
# many masks as there are classes. Let's ask the same query as above, but this
# time for *all* classes, not just the dog class: "For each pixel and each class
# C, is class C the most most likely class?"
#
# This one is a bit more involved, so we'll first show how to do it with a
# single image, and then we'll generalize to the batch
num_classes = normalized_masks.shape[1]
dog1_masks = normalized_masks[0]
def test_draw_segmentation_masks_errors():
h, w = 10, 10
masks = torch.randint(0, 2, size=(h, w), dtype=torch.bool)
img = torch.randint(0, 256, size=(3, h, w), dtype=torch.uint8)
with pytest.raises(TypeError, match="The image must be a tensor"):
utils.draw_segmentation_masks(image="Not A Tensor Image", masks=masks)
with pytest.raises(ValueError, match="The image dtype must be"):
img_bad_dtype = torch.randint(0,
256,
size=(3, h, w),
dtype=torch.int64)
utils.draw_segmentation_masks(image=img_bad_dtype, masks=masks)
with pytest.raises(ValueError,
match="Pass individual images, not batches"):
batch = torch.randint(0, 256, size=(10, 3, h, w), dtype=torch.uint8)
utils.draw_segmentation_masks(image=batch, masks=masks)
with pytest.raises(ValueError, match="Pass an RGB image"):
one_channel = torch.randint(0, 256, size=(1, h, w), dtype=torch.uint8)
utils.draw_segmentation_masks(image=one_channel, masks=masks)
with pytest.raises(ValueError, match="The masks must be of dtype bool"):
masks_bad_dtype = torch.randint(0, 2, size=(h, w), dtype=torch.float)
utils.draw_segmentation_masks(image=img, masks=masks_bad_dtype)
with pytest.raises(ValueError, match="masks must be of shape"):
masks_bad_shape = torch.randint(0,
2,
size=(3, 2, h, w),
dtype=torch.bool)
utils.draw_segmentation_masks(image=img, masks=masks_bad_shape)
with pytest.raises(ValueError,
match="must have the same height and width"):
masks_bad_shape = torch.randint(0,
2,
size=(h + 4, w),
dtype=torch.bool)
utils.draw_segmentation_masks(image=img, masks=masks_bad_shape)
with pytest.raises(ValueError, match="There are more masks"):
utils.draw_segmentation_masks(image=img, masks=masks, colors=[])
with pytest.raises(
ValueError,
match="colors must be a tuple or a string, or a list thereof"):
bad_colors = np.array(["red", "blue"]) # should be a list
utils.draw_segmentation_masks(image=img,
masks=masks,
colors=bad_colors)
with pytest.raises(
ValueError,
match="It seems that you passed a tuple of colors instead of"):
bad_colors = ("red", "blue") # should be a list
utils.draw_segmentation_masks(image=img,
masks=masks,
colors=bad_colors)
# transparency of masks.
#
# Here is demo with torchvision's FCN Resnet-50, loaded with
# :func:`~torchvision.models.segmentation.fcn_resnet50`.
# You can also try using
# DeepLabv3 (:func:`~torchvision.models.segmentation.deeplabv3_resnet50`)
# or lraspp mobilenet models
# (:func:`~torchvision.models.segmentation.lraspp_mobilenet_v3_large`).
#
# Like :func:`~torchvision.utils.draw_bounding_boxes`,
# :func:`~torchvision.utils.draw_segmentation_masks` requires a single RGB image
# of dtype `uint8`.
from torchvision.models.segmentation import fcn_resnet50
from torchvision.utils import draw_segmentation_masks
model = fcn_resnet50(pretrained=True, progress=False)
model = model.eval()
# The model expects the batch to be normalized
batch = F.normalize(batch,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
outputs = model(batch)
dogs_with_masks = [
draw_segmentation_masks(dog_int, masks=masks, alpha=0.6)
for dog_int, masks in zip((dog1_int, dog2_int), outputs['out'])
]
show(dogs_with_masks)
# For each instance, the boolean tensors represent if the particular pixel
# belongs to the segmentation mask of the image.
print(masks.size())
print(masks)
####################################
# Let us visualize an image and plot its corresponding segmentation masks.
# We will use the :func:`~torchvision.utils.draw_segmentation_masks` to draw the segmentation masks.
from torchvision.utils import draw_segmentation_masks
drawn_masks = []
for mask in masks:
drawn_masks.append(
draw_segmentation_masks(img, mask, alpha=0.8, colors="blue"))
show(drawn_masks)
####################################
# To convert the boolean masks into bounding boxes.
# We will use the :func:`~torchvision.ops.masks_to_boxes` from the torchvision.ops module
# It returns the boxes in ``(xmin, ymin, xmax, ymax)`` format.
from torchvision.ops import masks_to_boxes
boxes = masks_to_boxes(masks)
print(boxes.size())
print(boxes)
####################################
