def __getitem__(self, idx):
# load images and masks
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = read_image(img_path)
mask = read_image(mask_path)
img = F.convert_image_dtype(img, dtype=torch.float)
mask = F.convert_image_dtype(mask, dtype=torch.float)
# We get the unique colors, as these would be the object ids.
obj_ids = torch.unique(mask)
# first id is the background, so remove it.
obj_ids = obj_ids[1:]
# split the color-encoded mask into a set of boolean masks.
masks = mask == obj_ids[:, None, None]
boxes = masks_to_boxes(masks)
# there is only one class
labels = torch.ones((masks.shape[0], ), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def test_masks_to_boxes():
with PIL.Image.open(os.path.join(ASSETS_DIRECTORY, "masks.tiff")) as image:
masks = torch.zeros((image.n_frames, image.height, image.width),
dtype=torch.int)
for index in range(image.n_frames):
image.seek(index)
frame = numpy.array(image)
masks[index] = torch.tensor(frame)
expected = torch.tensor(
[[127, 2, 165, 40], [2, 50, 44, 92], [56, 63, 98, 100],
[139, 68, 175, 104], [160, 112, 198, 145], [49, 138, 99, 182],
[108, 148, 152, 213]],
dtype=torch.int32)
torch.testing.assert_close(masks_to_boxes(masks), expected)
def _copy_paste(
image: torch.Tensor,
target: Dict[str, Tensor],
paste_image: torch.Tensor,
paste_target: Dict[str, Tensor],
blending: bool = True,
resize_interpolation: F.InterpolationMode = F.InterpolationMode.BILINEAR,
) -> Tuple[torch.Tensor, Dict[str, Tensor]]:
# Random paste targets selection:
num_masks = len(paste_target["masks"])
if num_masks < 1:
# Such degerante case with num_masks=0 can happen with LSJ
# Let's just return (image, target)
return image, target
# We have to please torch script by explicitly specifying dtype as torch.long
random_selection = torch.randint(0,
num_masks, (num_masks, ),
device=paste_image.device)
random_selection = torch.unique(random_selection).to(torch.long)
paste_masks = paste_target["masks"][random_selection]
paste_boxes = paste_target["boxes"][random_selection]
paste_labels = paste_target["labels"][random_selection]
masks = target["masks"]
# We resize source and paste data if they have different sizes
# This is something we introduced here as originally the algorithm works
# on equal-sized data (for example, coming from LSJ data augmentations)
size1 = image.shape[-2:]
size2 = paste_image.shape[-2:]
if size1 != size2:
paste_image = F.resize(paste_image,
size1,
interpolation=resize_interpolation)
paste_masks = F.resize(paste_masks,
size1,
interpolation=F.InterpolationMode.NEAREST)
# resize bboxes:
ratios = torch.tensor((size1[1] / size2[1], size1[0] / size2[0]),
device=paste_boxes.device)
paste_boxes = paste_boxes.view(-1, 2,
2).mul(ratios).view(paste_boxes.shape)
paste_alpha_mask = paste_masks.sum(dim=0) > 0
if blending:
paste_alpha_mask = F.gaussian_blur(
paste_alpha_mask.unsqueeze(0),
kernel_size=(5, 5),
sigma=[
2.0,
],
)
# Copy-paste images:
image = (image * (~paste_alpha_mask)) + (paste_image * paste_alpha_mask)
# Copy-paste masks:
masks = masks * (~paste_alpha_mask)
non_all_zero_masks = masks.sum((-1, -2)) > 0
masks = masks[non_all_zero_masks]
# Do a shallow copy of the target dict
out_target = {k: v for k, v in target.items()}
out_target["masks"] = torch.cat([masks, paste_masks])
# Copy-paste boxes and labels
boxes = ops.masks_to_boxes(masks)
out_target["boxes"] = torch.cat([boxes, paste_boxes])
labels = target["labels"][non_all_zero_masks]
out_target["labels"] = torch.cat([labels, paste_labels])
# Update additional optional keys: area and iscrowd if exist
if "area" in target:
out_target["area"] = out_target["masks"].sum(
(-1, -2)).to(torch.float32)
if "iscrowd" in target and "iscrowd" in paste_target:
# target['iscrowd'] size can be differ from mask size (non_all_zero_masks)
# For example, if previous transforms geometrically modifies masks/boxes/labels but
# does not update "iscrowd"
if len(target["iscrowd"]) == len(non_all_zero_masks):
iscrowd = target["iscrowd"][non_all_zero_masks]
paste_iscrowd = paste_target["iscrowd"][random_selection]
out_target["iscrowd"] = torch.cat([iscrowd, paste_iscrowd])
# Check for degenerated boxes and remove them
boxes = out_target["boxes"]
degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]
if degenerate_boxes.any():
valid_targets = ~degenerate_boxes.any(dim=1)
out_target["boxes"] = boxes[valid_targets]
out_target["masks"] = out_target["masks"][valid_targets]
out_target["labels"] = out_target["labels"][valid_targets]
if "area" in out_target:
out_target["area"] = out_target["area"][valid_targets]
if "iscrowd" in out_target and len(
out_target["iscrowd"]) == len(valid_targets):
out_target["iscrowd"] = out_target["iscrowd"][valid_targets]
return image, out_target
# ----------------------------------------
# For example, the masks to bounding_boxes operation can be used to transform masks into bounding boxes that can be
# used in methods like Faster RCNN and YOLO.
with PIL.Image.open(os.path.join(ASSETS_DIRECTORY, "masks.tiff")) as image:
masks = torch.zeros((image.n_frames, image.height, image.width),
dtype=torch.int)
for index in range(image.n_frames):
image.seek(index)
frame = numpy.array(image)
masks[index] = torch.tensor(frame)
bounding_boxes = masks_to_boxes(masks)
figure = matplotlib.pyplot.figure()
a = figure.add_subplot(121)
b = figure.add_subplot(122)
labeled_image = torch.sum(masks, 0)
a.imshow(labeled_image)
b.imshow(labeled_image)
for bounding_box in bounding_boxes:
x0, y0, x1, y1 = bounding_box
rectangle = matplotlib.patches.Rectangle((x0, y0),
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)
####################################
# As the shape denotes, there are 3 boxes and in ``(xmin, ymin, xmax, ymax)`` format.
# These can be visualized very easily with :func:`~torchvision.utils.draw_bounding_boxes` utility
# provided in :ref:`torchvision.utils <utils>`.
from torchvision.utils import draw_bounding_boxes
drawn_boxes = draw_bounding_boxes(img, boxes, colors="red")
show(drawn_boxes)
###################################
# These boxes can now directly be used by detection models in torchvision.