OFA-Image_Caption / utils /transforms.py
JustinLin610's picture
update
75ba0e0
import random
import torch
import torchvision.transforms as T
import torchvision.transforms.functional as F
import numpy as np
from PIL import Image
def crop(image, target, region, delete=True):
cropped_image = F.crop(image, *region)
target = target.copy()
i, j, h, w = region
# should we do something wrt the original size?
target["size"] = torch.tensor([h, w])
fields = ["labels", "area"]
if "boxes" in target:
boxes = target["boxes"]
max_size = torch.as_tensor([w, h], dtype=torch.float32)
cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
cropped_boxes = cropped_boxes.clamp(min=0)
area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
target["boxes"] = cropped_boxes.reshape(-1, 4)
target["area"] = area
fields.append("boxes")
if "polygons" in target:
polygons = target["polygons"]
num_polygons = polygons.shape[0]
max_size = torch.as_tensor([w, h], dtype=torch.float32)
start_coord = torch.cat([torch.tensor([j, i], dtype=torch.float32)
for _ in range(polygons.shape[1] // 2)], dim=0)
cropped_boxes = polygons - start_coord
cropped_boxes = torch.min(cropped_boxes.reshape(num_polygons, -1, 2), max_size)
cropped_boxes = cropped_boxes.clamp(min=0)
target["polygons"] = cropped_boxes.reshape(num_polygons, -1)
fields.append("polygons")
if "masks" in target:
# FIXME should we update the area here if there are no boxes?
target['masks'] = target['masks'][:, i:i + h, j:j + w]
fields.append("masks")
# remove elements for which the boxes or masks that have zero area
if delete and ("boxes" in target or "masks" in target):
# favor boxes selection when defining which elements to keep
# this is compatible with previous implementation
if "boxes" in target:
cropped_boxes = target['boxes'].reshape(-1, 2, 2)
keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
else:
keep = target['masks'].flatten(1).any(1)
for field in fields:
target[field] = target[field][keep.tolist()]
return cropped_image, target
def hflip(image, target):
flipped_image = F.hflip(image)
w, h = image.size
target = target.copy()
if "boxes" in target:
boxes = target["boxes"]
boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])
target["boxes"] = boxes
if "polygons" in target:
polygons = target["polygons"]
num_polygons = polygons.shape[0]
polygons = polygons.reshape(num_polygons, -1, 2) * torch.as_tensor([-1, 1]) + torch.as_tensor([w, 0])
target["polygons"] = polygons
if "masks" in target:
target['masks'] = target['masks'].flip(-1)
return flipped_image, target
def resize(image, target, size, max_size=None):
# size can be min_size (scalar) or (w, h) tuple
def get_size_with_aspect_ratio(image_size, size, max_size=None):
w, h = image_size
if (w <= h and w == size) or (h <= w and h == size):
if max_size is not None:
max_size = int(max_size)
h = min(h, max_size)
w = min(w, max_size)
return (h, w)
if w < h:
ow = size
oh = int(size * h / w)
else:
oh = size
ow = int(size * w / h)
if max_size is not None:
max_size = int(max_size)
oh = min(oh, max_size)
ow = min(ow, max_size)
return (oh, ow)
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)
size = get_size(image.size, size, max_size)
rescaled_image = F.resize(image, size, interpolation=Image.BICUBIC)
if target is None:
return rescaled_image
ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
ratio_width, ratio_height = ratios
target = target.copy()
if "boxes" in target:
boxes = target["boxes"]
scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
target["boxes"] = scaled_boxes
if "polygons" in target:
polygons = target["polygons"]
scaled_ratio = torch.cat([torch.tensor([ratio_width, ratio_height])
for _ in range(polygons.shape[1] // 2)], dim=0)
scaled_polygons = polygons * scaled_ratio
target["polygons"] = scaled_polygons
if "area" in target:
area = target["area"]
scaled_area = area * (ratio_width * ratio_height)
target["area"] = scaled_area
h, w = size
target["size"] = torch.tensor([h, w])
if "masks" in target:
assert False
# target['masks'] = interpolate(
# target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
return rescaled_image, target
class CenterCrop(object):
def __init__(self, size):
self.size = size
def __call__(self, img, target):
image_width, image_height = img.size
crop_height, crop_width = self.size
crop_top = int(round((image_height - crop_height) / 2.))
crop_left = int(round((image_width - crop_width) / 2.))
return crop(img, target, (crop_top, crop_left, crop_height, crop_width))
class ObjectCenterCrop(object):
def __init__(self, size):
self.size = size
def __call__(self, img, target):
image_width, image_height = img.size
crop_height, crop_width = self.size
x0 = float(target['boxes'][0][0])
y0 = float(target['boxes'][0][1])
x1 = float(target['boxes'][0][2])
y1 = float(target['boxes'][0][3])
center_x = (x0 + x1) / 2
center_y = (y0 + y1) / 2
crop_left = max(center_x-crop_width/2 + min(image_width-center_x-crop_width/2, 0), 0)
crop_top = max(center_y-crop_height/2 + min(image_height-center_y-crop_height/2, 0), 0)
return crop(img, target, (crop_top, crop_left, crop_height, crop_width), delete=False)
class RandomHorizontalFlip(object):
def __init__(self, p=0.5):
self.p = p
def __call__(self, img, target):
if random.random() < self.p:
return hflip(img, target)
return img, target
class RandomResize(object):
def __init__(self, sizes, max_size=None, equal=False):
assert isinstance(sizes, (list, tuple))
self.sizes = sizes
self.max_size = max_size
self.equal = equal
def __call__(self, img, target=None):
size = random.choice(self.sizes)
if self.equal:
return resize(img, target, size, size)
else:
return resize(img, target, size, self.max_size)
class ToTensor(object):
def __call__(self, img, target):
return F.to_tensor(img), target
class Normalize(object):
def __init__(self, mean, std, max_image_size=512):
self.mean = mean
self.std = std
self.max_image_size = max_image_size
def __call__(self, image, target=None):
image = F.normalize(image, mean=self.mean, std=self.std)
if target is None:
return image, None
target = target.copy()
# h, w = image.shape[-2:]
h, w = target["size"][0], target["size"][1]
if "boxes" in target:
boxes = target["boxes"]
boxes = boxes / self.max_image_size
target["boxes"] = boxes
if "polygons" in target:
polygons = target["polygons"]
scale = torch.cat([torch.tensor([w, h], dtype=torch.float32)
for _ in range(polygons.shape[1] // 2)], dim=0)
polygons = polygons / scale
target["polygons"] = polygons
return image, target
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
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 LargeScaleJitter(object):
"""
implementation of large scale jitter from copy_paste
"""
def __init__(self, output_size=512, aug_scale_min=0.3, aug_scale_max=2.0):
self.desired_size = torch.tensor([output_size])
self.aug_scale_min = aug_scale_min
self.aug_scale_max = aug_scale_max
def rescale_target(self, scaled_size, image_size, target):
# compute rescaled targets
image_scale = scaled_size / image_size
ratio_height, ratio_width = image_scale
target = target.copy()
target["size"] = scaled_size
if "boxes" in target:
boxes = target["boxes"]
scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
target["boxes"] = scaled_boxes
if "area" in target:
area = target["area"]
scaled_area = area * (ratio_width * ratio_height)
target["area"] = scaled_area
if "masks" in target:
assert False
masks = target['masks']
# masks = interpolate(
# masks[:, None].float(), scaled_size, mode="nearest")[:, 0] > 0.5
target['masks'] = masks
return target
def crop_target(self, region, target):
i, j, h, w = region
fields = ["labels", "area"]
target = target.copy()
target["size"] = torch.tensor([h, w])
if "boxes" in target:
boxes = target["boxes"]
max_size = torch.as_tensor([w, h], dtype=torch.float32)
cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
cropped_boxes = cropped_boxes.clamp(min=0)
area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
target["boxes"] = cropped_boxes.reshape(-1, 4)
target["area"] = area
fields.append("boxes")
if "masks" in target:
# FIXME should we update the area here if there are no boxes?
target['masks'] = target['masks'][:, i:i + h, j:j + w]
fields.append("masks")
# remove elements for which the boxes or masks that have zero area
if "boxes" in target or "masks" in target:
# favor boxes selection when defining which elements to keep
# this is compatible with previous implementation
if "boxes" in target:
cropped_boxes = target['boxes'].reshape(-1, 2, 2)
keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
else:
keep = target['masks'].flatten(1).any(1)
for field in fields:
target[field] = target[field][keep.tolist()]
return target
def pad_target(self, padding, target):
target = target.copy()
if "masks" in target:
target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[1], 0, padding[0]))
return target
def __call__(self, image, target=None):
image_size = image.size
image_size = torch.tensor(image_size[::-1])
random_scale = torch.rand(1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min
scaled_size = (random_scale * self.desired_size).round()
scale = torch.maximum(scaled_size / image_size[0], scaled_size / image_size[1])
scaled_size = (image_size * scale).round().int()
scaled_image = F.resize(image, scaled_size.tolist(), interpolation=Image.BICUBIC)
if target is not None:
target = self.rescale_target(scaled_size, image_size, target)
# randomly crop or pad images
if random_scale >= 1:
# Selects non-zero random offset (x, y) if scaled image is larger than desired_size.
max_offset = scaled_size - self.desired_size
offset = (max_offset * torch.rand(2)).floor().int()
region = (offset[0].item(), offset[1].item(),
self.desired_size[0].item(), self.desired_size[0].item())
output_image = F.crop(scaled_image, *region)
if target is not None:
target = self.crop_target(region, target)
else:
assert False
padding = self.desired_size - scaled_size
output_image = F.pad(scaled_image, [0, 0, padding[1].item(), padding[0].item()])
if target is not None:
target = self.pad_target(padding, target)
return output_image, target
class OriginLargeScaleJitter(object):
"""
implementation of large scale jitter from copy_paste
"""
def __init__(self, output_size=512, aug_scale_min=0.3, aug_scale_max=2.0):
self.desired_size = torch.tensor(output_size)
self.aug_scale_min = aug_scale_min
self.aug_scale_max = aug_scale_max
def rescale_target(self, scaled_size, image_size, target):
# compute rescaled targets
image_scale = scaled_size / image_size
ratio_height, ratio_width = image_scale
target = target.copy()
target["size"] = scaled_size
if "boxes" in target:
boxes = target["boxes"]
scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
target["boxes"] = scaled_boxes
if "area" in target:
area = target["area"]
scaled_area = area * (ratio_width * ratio_height)
target["area"] = scaled_area
if "masks" in target:
assert False
masks = target['masks']
# masks = interpolate(
# masks[:, None].float(), scaled_size, mode="nearest")[:, 0] > 0.5
target['masks'] = masks
return target
def crop_target(self, region, target):
i, j, h, w = region
fields = ["labels", "area"]
target = target.copy()
target["size"] = torch.tensor([h, w])
if "boxes" in target:
boxes = target["boxes"]
max_size = torch.as_tensor([w, h], dtype=torch.float32)
cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
cropped_boxes = cropped_boxes.clamp(min=0)
area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
target["boxes"] = cropped_boxes.reshape(-1, 4)
target["area"] = area
fields.append("boxes")
if "masks" in target:
# FIXME should we update the area here if there are no boxes?
target['masks'] = target['masks'][:, i:i + h, j:j + w]
fields.append("masks")
# remove elements for which the boxes or masks that have zero area
if "boxes" in target or "masks" in target:
# favor boxes selection when defining which elements to keep
# this is compatible with previous implementation
if "boxes" in target:
cropped_boxes = target['boxes'].reshape(-1, 2, 2)
keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
else:
keep = target['masks'].flatten(1).any(1)
for field in fields:
target[field] = target[field][keep.tolist()]
return target
def pad_target(self, padding, target):
target = target.copy()
if "masks" in target:
target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[1], 0, padding[0]))
return target
def __call__(self, image, target=None):
image_size = image.size
image_size = torch.tensor(image_size[::-1])
out_desired_size = (self.desired_size * image_size / max(image_size)).round().int()
random_scale = torch.rand(1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min
scaled_size = (random_scale * self.desired_size).round()
scale = torch.minimum(scaled_size / image_size[0], scaled_size / image_size[1])
scaled_size = (image_size * scale).round().int()
scaled_image = F.resize(image, scaled_size.tolist())
if target is not None:
target = self.rescale_target(scaled_size, image_size, target)
# randomly crop or pad images
if random_scale > 1:
# Selects non-zero random offset (x, y) if scaled image is larger than desired_size.
max_offset = scaled_size - out_desired_size
offset = (max_offset * torch.rand(2)).floor().int()
region = (offset[0].item(), offset[1].item(),
out_desired_size[0].item(), out_desired_size[1].item())
output_image = F.crop(scaled_image, *region)
if target is not None:
target = self.crop_target(region, target)
else:
padding = out_desired_size - scaled_size
output_image = F.pad(scaled_image, [0, 0, padding[1].item(), padding[0].item()])
if target is not None:
target = self.pad_target(padding, target)
return output_image, target
class RandomDistortion(object):
"""
Distort image w.r.t hue, saturation and exposure.
"""
def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, prob=0.5):
self.prob = prob
self.tfm = T.ColorJitter(brightness, contrast, saturation, hue)
def __call__(self, img, target=None):
if np.random.random() < self.prob:
return self.tfm(img), target
else:
return img, target