529 lines
18 KiB
Python
529 lines
18 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""
|
|
Transforms and data augmentation for both image + bbox.
|
|
"""
|
|
|
|
import logging
|
|
|
|
import random
|
|
from typing import Iterable
|
|
|
|
import torch
|
|
import torchvision.transforms as T
|
|
import torchvision.transforms.functional as F
|
|
import torchvision.transforms.v2.functional as Fv2
|
|
from PIL import Image as PILImage
|
|
|
|
from torchvision.transforms import InterpolationMode
|
|
|
|
from training.utils.data_utils import VideoDatapoint
|
|
|
|
|
|
def hflip(datapoint, index):
|
|
|
|
datapoint.frames[index].data = F.hflip(datapoint.frames[index].data)
|
|
for obj in datapoint.frames[index].objects:
|
|
if obj.segment is not None:
|
|
obj.segment = F.hflip(obj.segment)
|
|
|
|
return datapoint
|
|
|
|
|
|
def get_size_with_aspect_ratio(image_size, size, max_size=None):
|
|
w, h = image_size
|
|
if max_size is not None:
|
|
min_original_size = float(min((w, h)))
|
|
max_original_size = float(max((w, h)))
|
|
if max_original_size / min_original_size * size > max_size:
|
|
size = max_size * min_original_size / max_original_size
|
|
|
|
if (w <= h and w == size) or (h <= w and h == size):
|
|
return (h, w)
|
|
|
|
if w < h:
|
|
ow = int(round(size))
|
|
oh = int(round(size * h / w))
|
|
else:
|
|
oh = int(round(size))
|
|
ow = int(round(size * w / h))
|
|
|
|
return (oh, ow)
|
|
|
|
|
|
def resize(datapoint, index, size, max_size=None, square=False, v2=False):
|
|
# size can be min_size (scalar) or (w, h) tuple
|
|
|
|
def get_size(image_size, size, max_size=None):
|
|
if isinstance(size, (list, tuple)):
|
|
return size[::-1]
|
|
else:
|
|
return get_size_with_aspect_ratio(image_size, size, max_size)
|
|
|
|
if square:
|
|
size = size, size
|
|
else:
|
|
cur_size = (
|
|
datapoint.frames[index].data.size()[-2:][::-1]
|
|
if v2
|
|
else datapoint.frames[index].data.size
|
|
)
|
|
size = get_size(cur_size, size, max_size)
|
|
|
|
old_size = (
|
|
datapoint.frames[index].data.size()[-2:][::-1]
|
|
if v2
|
|
else datapoint.frames[index].data.size
|
|
)
|
|
if v2:
|
|
datapoint.frames[index].data = Fv2.resize(
|
|
datapoint.frames[index].data, size, antialias=True
|
|
)
|
|
else:
|
|
datapoint.frames[index].data = F.resize(datapoint.frames[index].data, size)
|
|
|
|
new_size = (
|
|
datapoint.frames[index].data.size()[-2:][::-1]
|
|
if v2
|
|
else datapoint.frames[index].data.size
|
|
)
|
|
|
|
for obj in datapoint.frames[index].objects:
|
|
if obj.segment is not None:
|
|
obj.segment = F.resize(obj.segment[None, None], size).squeeze()
|
|
|
|
h, w = size
|
|
datapoint.frames[index].size = (h, w)
|
|
return datapoint
|
|
|
|
|
|
def pad(datapoint, index, padding, v2=False):
|
|
old_h, old_w = datapoint.frames[index].size
|
|
h, w = old_h, old_w
|
|
if len(padding) == 2:
|
|
# assumes that we only pad on the bottom right corners
|
|
datapoint.frames[index].data = F.pad(
|
|
datapoint.frames[index].data, (0, 0, padding[0], padding[1])
|
|
)
|
|
h += padding[1]
|
|
w += padding[0]
|
|
else:
|
|
# left, top, right, bottom
|
|
datapoint.frames[index].data = F.pad(
|
|
datapoint.frames[index].data,
|
|
(padding[0], padding[1], padding[2], padding[3]),
|
|
)
|
|
h += padding[1] + padding[3]
|
|
w += padding[0] + padding[2]
|
|
|
|
datapoint.frames[index].size = (h, w)
|
|
|
|
for obj in datapoint.frames[index].objects:
|
|
if obj.segment is not None:
|
|
if v2:
|
|
if len(padding) == 2:
|
|
obj.segment = Fv2.pad(obj.segment, (0, 0, padding[0], padding[1]))
|
|
else:
|
|
obj.segment = Fv2.pad(obj.segment, tuple(padding))
|
|
else:
|
|
if len(padding) == 2:
|
|
obj.segment = F.pad(obj.segment, (0, 0, padding[0], padding[1]))
|
|
else:
|
|
obj.segment = F.pad(obj.segment, tuple(padding))
|
|
return datapoint
|
|
|
|
|
|
class RandomHorizontalFlip:
|
|
def __init__(self, consistent_transform, p=0.5):
|
|
self.p = p
|
|
self.consistent_transform = consistent_transform
|
|
|
|
def __call__(self, datapoint, **kwargs):
|
|
if self.consistent_transform:
|
|
if random.random() < self.p:
|
|
for i in range(len(datapoint.frames)):
|
|
datapoint = hflip(datapoint, i)
|
|
return datapoint
|
|
for i in range(len(datapoint.frames)):
|
|
if random.random() < self.p:
|
|
datapoint = hflip(datapoint, i)
|
|
return datapoint
|
|
|
|
|
|
class RandomResizeAPI:
|
|
def __init__(
|
|
self, sizes, consistent_transform, max_size=None, square=False, v2=False
|
|
):
|
|
if isinstance(sizes, int):
|
|
sizes = (sizes,)
|
|
assert isinstance(sizes, Iterable)
|
|
self.sizes = list(sizes)
|
|
self.max_size = max_size
|
|
self.square = square
|
|
self.consistent_transform = consistent_transform
|
|
self.v2 = v2
|
|
|
|
def __call__(self, datapoint, **kwargs):
|
|
if self.consistent_transform:
|
|
size = random.choice(self.sizes)
|
|
for i in range(len(datapoint.frames)):
|
|
datapoint = resize(
|
|
datapoint, i, size, self.max_size, square=self.square, v2=self.v2
|
|
)
|
|
return datapoint
|
|
for i in range(len(datapoint.frames)):
|
|
size = random.choice(self.sizes)
|
|
datapoint = resize(
|
|
datapoint, i, size, self.max_size, square=self.square, v2=self.v2
|
|
)
|
|
return datapoint
|
|
|
|
|
|
class ToTensorAPI:
|
|
def __init__(self, v2=False):
|
|
self.v2 = v2
|
|
|
|
def __call__(self, datapoint: VideoDatapoint, **kwargs):
|
|
for img in datapoint.frames:
|
|
if self.v2:
|
|
img.data = Fv2.to_image_tensor(img.data)
|
|
else:
|
|
img.data = F.to_tensor(img.data)
|
|
return datapoint
|
|
|
|
|
|
class NormalizeAPI:
|
|
def __init__(self, mean, std, v2=False):
|
|
self.mean = mean
|
|
self.std = std
|
|
self.v2 = v2
|
|
|
|
def __call__(self, datapoint: VideoDatapoint, **kwargs):
|
|
for img in datapoint.frames:
|
|
if self.v2:
|
|
img.data = Fv2.convert_image_dtype(img.data, torch.float32)
|
|
img.data = Fv2.normalize(img.data, mean=self.mean, std=self.std)
|
|
else:
|
|
img.data = F.normalize(img.data, mean=self.mean, std=self.std)
|
|
|
|
return datapoint
|
|
|
|
|
|
class ComposeAPI:
|
|
def __init__(self, transforms):
|
|
self.transforms = transforms
|
|
|
|
def __call__(self, datapoint, **kwargs):
|
|
for t in self.transforms:
|
|
datapoint = t(datapoint, **kwargs)
|
|
return datapoint
|
|
|
|
def __repr__(self):
|
|
format_string = self.__class__.__name__ + "("
|
|
for t in self.transforms:
|
|
format_string += "\n"
|
|
format_string += " {0}".format(t)
|
|
format_string += "\n)"
|
|
return format_string
|
|
|
|
|
|
class RandomGrayscale:
|
|
def __init__(self, consistent_transform, p=0.5):
|
|
self.p = p
|
|
self.consistent_transform = consistent_transform
|
|
self.Grayscale = T.Grayscale(num_output_channels=3)
|
|
|
|
def __call__(self, datapoint: VideoDatapoint, **kwargs):
|
|
if self.consistent_transform:
|
|
if random.random() < self.p:
|
|
for img in datapoint.frames:
|
|
img.data = self.Grayscale(img.data)
|
|
return datapoint
|
|
for img in datapoint.frames:
|
|
if random.random() < self.p:
|
|
img.data = self.Grayscale(img.data)
|
|
return datapoint
|
|
|
|
|
|
class ColorJitter:
|
|
def __init__(self, consistent_transform, brightness, contrast, saturation, hue):
|
|
self.consistent_transform = consistent_transform
|
|
self.brightness = (
|
|
brightness
|
|
if isinstance(brightness, list)
|
|
else [max(0, 1 - brightness), 1 + brightness]
|
|
)
|
|
self.contrast = (
|
|
contrast
|
|
if isinstance(contrast, list)
|
|
else [max(0, 1 - contrast), 1 + contrast]
|
|
)
|
|
self.saturation = (
|
|
saturation
|
|
if isinstance(saturation, list)
|
|
else [max(0, 1 - saturation), 1 + saturation]
|
|
)
|
|
self.hue = hue if isinstance(hue, list) or hue is None else ([-hue, hue])
|
|
|
|
def __call__(self, datapoint: VideoDatapoint, **kwargs):
|
|
if self.consistent_transform:
|
|
# Create a color jitter transformation params
|
|
(
|
|
fn_idx,
|
|
brightness_factor,
|
|
contrast_factor,
|
|
saturation_factor,
|
|
hue_factor,
|
|
) = T.ColorJitter.get_params(
|
|
self.brightness, self.contrast, self.saturation, self.hue
|
|
)
|
|
for img in datapoint.frames:
|
|
if not self.consistent_transform:
|
|
(
|
|
fn_idx,
|
|
brightness_factor,
|
|
contrast_factor,
|
|
saturation_factor,
|
|
hue_factor,
|
|
) = T.ColorJitter.get_params(
|
|
self.brightness, self.contrast, self.saturation, self.hue
|
|
)
|
|
for fn_id in fn_idx:
|
|
if fn_id == 0 and brightness_factor is not None:
|
|
img.data = F.adjust_brightness(img.data, brightness_factor)
|
|
elif fn_id == 1 and contrast_factor is not None:
|
|
img.data = F.adjust_contrast(img.data, contrast_factor)
|
|
elif fn_id == 2 and saturation_factor is not None:
|
|
img.data = F.adjust_saturation(img.data, saturation_factor)
|
|
elif fn_id == 3 and hue_factor is not None:
|
|
img.data = F.adjust_hue(img.data, hue_factor)
|
|
return datapoint
|
|
|
|
|
|
class RandomAffine:
|
|
def __init__(
|
|
self,
|
|
degrees,
|
|
consistent_transform,
|
|
scale=None,
|
|
translate=None,
|
|
shear=None,
|
|
image_mean=(123, 116, 103),
|
|
log_warning=True,
|
|
num_tentatives=1,
|
|
image_interpolation="bicubic",
|
|
):
|
|
"""
|
|
The mask is required for this transform.
|
|
if consistent_transform if True, then the same random affine is applied to all frames and masks.
|
|
"""
|
|
self.degrees = degrees if isinstance(degrees, list) else ([-degrees, degrees])
|
|
self.scale = scale
|
|
self.shear = (
|
|
shear if isinstance(shear, list) else ([-shear, shear] if shear else None)
|
|
)
|
|
self.translate = translate
|
|
self.fill_img = image_mean
|
|
self.consistent_transform = consistent_transform
|
|
self.log_warning = log_warning
|
|
self.num_tentatives = num_tentatives
|
|
|
|
if image_interpolation == "bicubic":
|
|
self.image_interpolation = InterpolationMode.BICUBIC
|
|
elif image_interpolation == "bilinear":
|
|
self.image_interpolation = InterpolationMode.BILINEAR
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def __call__(self, datapoint: VideoDatapoint, **kwargs):
|
|
for _tentative in range(self.num_tentatives):
|
|
res = self.transform_datapoint(datapoint)
|
|
if res is not None:
|
|
return res
|
|
|
|
if self.log_warning:
|
|
logging.warning(
|
|
f"Skip RandomAffine for zero-area mask in first frame after {self.num_tentatives} tentatives"
|
|
)
|
|
return datapoint
|
|
|
|
def transform_datapoint(self, datapoint: VideoDatapoint):
|
|
_, height, width = F.get_dimensions(datapoint.frames[0].data)
|
|
img_size = [width, height]
|
|
|
|
if self.consistent_transform:
|
|
# Create a random affine transformation
|
|
affine_params = T.RandomAffine.get_params(
|
|
degrees=self.degrees,
|
|
translate=self.translate,
|
|
scale_ranges=self.scale,
|
|
shears=self.shear,
|
|
img_size=img_size,
|
|
)
|
|
|
|
for img_idx, img in enumerate(datapoint.frames):
|
|
this_masks = [
|
|
obj.segment.unsqueeze(0) if obj.segment is not None else None
|
|
for obj in img.objects
|
|
]
|
|
if not self.consistent_transform:
|
|
# if not consistent we create a new affine params for every frame&mask pair Create a random affine transformation
|
|
affine_params = T.RandomAffine.get_params(
|
|
degrees=self.degrees,
|
|
translate=self.translate,
|
|
scale_ranges=self.scale,
|
|
shears=self.shear,
|
|
img_size=img_size,
|
|
)
|
|
|
|
transformed_bboxes, transformed_masks = [], []
|
|
for i in range(len(img.objects)):
|
|
if this_masks[i] is None:
|
|
transformed_masks.append(None)
|
|
# Dummy bbox for a dummy target
|
|
transformed_bboxes.append(torch.tensor([[0, 0, 1, 1]]))
|
|
else:
|
|
transformed_mask = F.affine(
|
|
this_masks[i],
|
|
*affine_params,
|
|
interpolation=InterpolationMode.NEAREST,
|
|
fill=0.0,
|
|
)
|
|
if img_idx == 0 and transformed_mask.max() == 0:
|
|
# We are dealing with a video and the object is not visible in the first frame
|
|
# Return the datapoint without transformation
|
|
return None
|
|
transformed_masks.append(transformed_mask.squeeze())
|
|
|
|
for i in range(len(img.objects)):
|
|
img.objects[i].segment = transformed_masks[i]
|
|
|
|
img.data = F.affine(
|
|
img.data,
|
|
*affine_params,
|
|
interpolation=self.image_interpolation,
|
|
fill=self.fill_img,
|
|
)
|
|
return datapoint
|
|
|
|
|
|
def random_mosaic_frame(
|
|
datapoint,
|
|
index,
|
|
grid_h,
|
|
grid_w,
|
|
target_grid_y,
|
|
target_grid_x,
|
|
should_hflip,
|
|
):
|
|
# Step 1: downsize the images and paste them into a mosaic
|
|
image_data = datapoint.frames[index].data
|
|
is_pil = isinstance(image_data, PILImage.Image)
|
|
if is_pil:
|
|
H_im = image_data.height
|
|
W_im = image_data.width
|
|
image_data_output = PILImage.new("RGB", (W_im, H_im))
|
|
else:
|
|
H_im = image_data.size(-2)
|
|
W_im = image_data.size(-1)
|
|
image_data_output = torch.zeros_like(image_data)
|
|
|
|
downsize_cache = {}
|
|
for grid_y in range(grid_h):
|
|
for grid_x in range(grid_w):
|
|
y_offset_b = grid_y * H_im // grid_h
|
|
x_offset_b = grid_x * W_im // grid_w
|
|
y_offset_e = (grid_y + 1) * H_im // grid_h
|
|
x_offset_e = (grid_x + 1) * W_im // grid_w
|
|
H_im_downsize = y_offset_e - y_offset_b
|
|
W_im_downsize = x_offset_e - x_offset_b
|
|
|
|
if (H_im_downsize, W_im_downsize) in downsize_cache:
|
|
image_data_downsize = downsize_cache[(H_im_downsize, W_im_downsize)]
|
|
else:
|
|
image_data_downsize = F.resize(
|
|
image_data,
|
|
size=(H_im_downsize, W_im_downsize),
|
|
interpolation=InterpolationMode.BILINEAR,
|
|
antialias=True, # antialiasing for downsizing
|
|
)
|
|
downsize_cache[(H_im_downsize, W_im_downsize)] = image_data_downsize
|
|
if should_hflip[grid_y, grid_x].item():
|
|
image_data_downsize = F.hflip(image_data_downsize)
|
|
|
|
if is_pil:
|
|
image_data_output.paste(image_data_downsize, (x_offset_b, y_offset_b))
|
|
else:
|
|
image_data_output[:, y_offset_b:y_offset_e, x_offset_b:x_offset_e] = (
|
|
image_data_downsize
|
|
)
|
|
|
|
datapoint.frames[index].data = image_data_output
|
|
|
|
# Step 2: downsize the masks and paste them into the target grid of the mosaic
|
|
for obj in datapoint.frames[index].objects:
|
|
if obj.segment is None:
|
|
continue
|
|
assert obj.segment.shape == (H_im, W_im) and obj.segment.dtype == torch.uint8
|
|
segment_output = torch.zeros_like(obj.segment)
|
|
|
|
target_y_offset_b = target_grid_y * H_im // grid_h
|
|
target_x_offset_b = target_grid_x * W_im // grid_w
|
|
target_y_offset_e = (target_grid_y + 1) * H_im // grid_h
|
|
target_x_offset_e = (target_grid_x + 1) * W_im // grid_w
|
|
target_H_im_downsize = target_y_offset_e - target_y_offset_b
|
|
target_W_im_downsize = target_x_offset_e - target_x_offset_b
|
|
|
|
segment_downsize = F.resize(
|
|
obj.segment[None, None],
|
|
size=(target_H_im_downsize, target_W_im_downsize),
|
|
interpolation=InterpolationMode.BILINEAR,
|
|
antialias=True, # antialiasing for downsizing
|
|
)[0, 0]
|
|
if should_hflip[target_grid_y, target_grid_x].item():
|
|
segment_downsize = F.hflip(segment_downsize[None, None])[0, 0]
|
|
|
|
segment_output[
|
|
target_y_offset_b:target_y_offset_e, target_x_offset_b:target_x_offset_e
|
|
] = segment_downsize
|
|
obj.segment = segment_output
|
|
|
|
return datapoint
|
|
|
|
|
|
class RandomMosaicVideoAPI:
|
|
def __init__(self, prob=0.15, grid_h=2, grid_w=2, use_random_hflip=False):
|
|
self.prob = prob
|
|
self.grid_h = grid_h
|
|
self.grid_w = grid_w
|
|
self.use_random_hflip = use_random_hflip
|
|
|
|
def __call__(self, datapoint, **kwargs):
|
|
if random.random() > self.prob:
|
|
return datapoint
|
|
|
|
# select a random location to place the target mask in the mosaic
|
|
target_grid_y = random.randint(0, self.grid_h - 1)
|
|
target_grid_x = random.randint(0, self.grid_w - 1)
|
|
# whether to flip each grid in the mosaic horizontally
|
|
if self.use_random_hflip:
|
|
should_hflip = torch.rand(self.grid_h, self.grid_w) < 0.5
|
|
else:
|
|
should_hflip = torch.zeros(self.grid_h, self.grid_w, dtype=torch.bool)
|
|
for i in range(len(datapoint.frames)):
|
|
datapoint = random_mosaic_frame(
|
|
datapoint,
|
|
i,
|
|
grid_h=self.grid_h,
|
|
grid_w=self.grid_w,
|
|
target_grid_y=target_grid_y,
|
|
target_grid_x=target_grid_x,
|
|
should_hflip=should_hflip,
|
|
)
|
|
|
|
return datapoint
|