def __call__(self, img):
"""
Args:
img (Tensor): Tensor image of size (C, H, W) to be erased.
Returns:
img (Tensor): Erased Tensor image.
"""
if self.same_place:
if self.random_p < self.p:
if self.randomize:
self.erase_params = self.get_params(img,
scale=self.scale,
ratio=self.ratio,
value=self.value)
self.randomize = False
x, y, h, w, v = self.erase_params
return F.erase(img, x, y, h, w, v, self.inplace)
return img
else:
self.random_p = random.random()
if self.random_p < self.p:
x, y, h, w, v = self.get_params(img,
scale=self.scale,
ratio=self.ratio,
value=self.value)
return F.erase(img, x, y, h, w, v, self.inplace)
return img
def __call__(self, sample):
image, label, add_data = sample
if random.random() < self.t.p:
params = self.t.get_params(image, self.t.scale, self.t.ratio,
self.t.value)
image = Fv.erase(image, *params, self.t.inplace)
if self.label_tf:
label = Fv.erase(label, *params, self.t.inplace)
if add_data is not None:
for i in range(len(add_data)):
add_data[i] = Fv.erase(add_data[i], *params,
self.t.inplace)
return (image, label, add_data)
def crop_out(
x: torch.Tensor,
interval: int,
magnitude: float,
max_magnitude: int,
fill: float,
) -> torch.Tensor:
# From: https://github.com/tensorflow/tpu/blob/3679ca6b979349dde6da7156be2528428b000c7c/models/official/efficientnet/autoaugment.py#L125
H, W = x.shape[-2], x.shape[-1]
# From: https://github.com/tensorflow/tpu/blob/3679ca6b979349dde6da7156be2528428b000c7c/models/official/efficientnet/autoaugment.py#L525
# From: https://github.com/tensorflow/tpu/blob/3679ca6b979349dde6da7156be2528428b000c7c/models/official/efficientnet/autoaugment.py#L666-L667
lateral_len = int(
(magnitude / max_magnitude) * interval) # Default: interval = 100
x_coor = torch.randint(0, W - lateral_len, []).item()
y_coor = torch.randint(0, H - lateral_len, []).item()
return TF.erase(x,
i=x_coor,
j=y_coor,
h=lateral_len,
w=lateral_len,
v=fill)
def forward(self, img):
if isinstance(img, Image.Image):
img = F.to_tensor(img)
i, j, erase_h, erase_w, value = self.get_params(img=img,
scale=self.scale,
value=self.value)
return TF.erase(img, i, j, erase_h, erase_w, value, True)
def __call__(self, img, target):
if random.uniform(0, 1) < self.p:
img = F.to_tensor(img)
x, y, h, w, v = self.eraser.get_params(
img, scale=self.scale, ratio=self.ratio, value=self.value)
img = F.erase(img, x, y, h, w, v, self.inplace)
img = F.to_pil_image(img)
# target
fields = ['boxes', "labels", "area", "iscrowd", "ignore", "track_ids"]
if 'boxes' in target:
erased_box = torch.tensor([[y, x, y + w, x + h]]).float()
lt = torch.max(erased_box[:, None, :2], target['boxes'][:, :2]) # [N,M,2]
rb = torch.min(erased_box[:, None, 2:], target['boxes'][:, 2:]) # [N,M,2]
wh = (rb - lt).clamp(min=0) # [N,M,2]
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
keep = inter[0] <= 0.7 * target['area']
left = torch.logical_and(
target['boxes'][:, 0] < erased_box[:, 0],
target['boxes'][:, 2] > erased_box[:, 0])
left = torch.logical_and(left, inter[0].bool())
right = torch.logical_and(
target['boxes'][:, 0] < erased_box[:, 2],
target['boxes'][:, 2] > erased_box[:, 2])
right = torch.logical_and(right, inter[0].bool())
top = torch.logical_and(
target['boxes'][:, 1] < erased_box[:, 1],
target['boxes'][:, 3] > erased_box[:, 1])
top = torch.logical_and(top, inter[0].bool())
bottom = torch.logical_and(
target['boxes'][:, 1] < erased_box[:, 3],
target['boxes'][:, 3] > erased_box[:, 3])
bottom = torch.logical_and(bottom, inter[0].bool())
only_one_crop = (top.float() + bottom.float() + left.float() + right.float()) > 1
left[only_one_crop] = False
right[only_one_crop] = False
top[only_one_crop] = False
bottom[only_one_crop] = False
target['boxes'][:, 2][left] = erased_box[:, 0]
target['boxes'][:, 0][right] = erased_box[:, 2]
target['boxes'][:, 3][top] = erased_box[:, 1]
target['boxes'][:, 1][bottom] = erased_box[:, 3]
for field in fields:
if field in target:
target[field] = target[field][keep]
return img, target
def __call__(self, img):
if random.uniform(0, 1) < self.p:
x, y, h, w, v = self.get_params(img,
scale=self.scale,
ratio=self.ratio,
value=self.value)
return tf.erase(img, x, y, h, w, v, self.inplace)
return img
def forward(self, img1, img2, flow, valid_flow_mask):
if torch.rand(1) > self.p:
return img1, img2, flow, valid_flow_mask
for _ in range(torch.randint(self.max_erase, size=(1,)).item()):
x, y, h, w, v = self.get_params(img2, scale=self.scale, ratio=self.ratio, value=[self.value])
img2 = F.erase(img2, x, y, h, w, v, self.inplace)
return img1, img2, flow, valid_flow_mask
def __call__(self, img):
x, y, h, w, v = self.get_params(img,
scale=self.scale,
ratio=self.ratio,
value=int(self.value))
mask = torch.ones_like(img).float()
mask[:, x:x + h, y:y + w] = 0.0
return F.erase(img, x, y, h, w, v, self.inplace), \
ToTensor()(ToPILImage()(mask)),\
ToTensor()(F.crop(ToPILImage()(img), x, y, h, w))
def __get_mask_batch(self):
mask_batch = []
indices = []
for i in torch.arange(0, self.image_tensor.shape[1], self.mask_size):
for j in torch.arange(0, self.image_tensor.shape[2], self.mask_size):
mask_batch.append(TF.erase(self.image_tensor, i, j, self.mask_size, self.mask_size, 0.))
indices.append((slice(i,i+self.mask_size), slice(j,j+self.mask_size)))
if len(mask_batch) == self.batch_size:
yield torch.stack(mask_batch), indices
mask_batch, indices = [], []
def __call__(self, image, label):
"""
Args:
img (Tensor): Tensor image of size (C, H, W) to be erased.
Returns:
img (Tensor): Erased Tensor image.
"""
if random.uniform(0, 1) < self.p:
x, y, h, w, v = self.get_params(image,
scale=self.scale,
ratio=self.ratio,
value=self.value)
image = TF.erase(image, x, y, h, w, v, self.inplace)
if self.erase_label:
label = label.unsqueeze(0)
label = TF.erase(label, x, y, h, w, v, self.inplace)
label = label.squeeze(0)
return image, label
def __call__(self, img):
img = F.to_tensor(img)
H, W = img.size()[-2:]
lui = torch.randint(0, H, (1, )).item()
luj = torch.randint(0, W, (1, )).item()
rbi = torch.randint(lui, H, (1, )).item()
rbj = torch.randint(luj, W, (1, )).item()
h = rbj - luj
w = rbi - lui
return F.to_pil_image(F.erase(img, lui, luj, h, w, 0))
def __call__(self, img, target=None):
img = F.to_tensor(img)
img_c, img_w, img_h = img.shape
rescaled_boxes = []
rescaled_box_images = []
for box in target['boxes']:
y1, x1, y2, x2 = box.int().tolist()
w = x2 - x1
h = y2 - y1
box_img = img[:, x1:x2, y1:y2]
random_scale = random.uniform(0.5, 2.0)
scaled_width = int(random_scale * w)
scaled_height = int(random_scale * h)
box_img = F.to_pil_image(box_img)
rescaled_box_image = F.resize(
box_img,
(scaled_width, scaled_height))
rescaled_box_images.append(F.to_tensor(rescaled_box_image))
rescaled_boxes.append([y1, x1, y1 + scaled_height, x1 + scaled_width])
for box in target['boxes']:
y1, x1, y2, x2 = box.int().tolist()
w = x2 - x1
h = y2 - y1
erase_value = torch.empty(
[img_c, w, h],
dtype=torch.float32).normal_()
img = F.erase(
img, x1, y1, w, h, erase_value, True)
for box, rescaled_box_image in zip(target['boxes'], rescaled_box_images):
y1, x1, y2, x2 = box.int().tolist()
w = x2 - x1
h = y2 - y1
_, scaled_width, scaled_height = rescaled_box_image.shape
rescaled_box_image = rescaled_box_image[
:,
:scaled_width - max(x1 + scaled_width - img_w, 0),
:scaled_height - max(y1 + scaled_height - img_h, 0)]
img[:, x1:x1 + scaled_width, y1:y1 + scaled_height] = rescaled_box_image
target['boxes'] = torch.tensor(rescaled_boxes).float()
img = F.to_pil_image(img)
return img, target
def __call__(self, img):
"""
Args:
img (Tensor): Tensor image of size (C, H, W) to be erased.
Returns:
img (Tensor): Erased Tensor image.
"""
if random.uniform(0, 1) < self.p:
x, y, h, w, v = self.get_params(img,
scale=self.scale,
ratio=self.ratio,
value=self.value)
return F.erase(img, x, y, h, w, v, self.inplace)
return img
def _tensor_cutout(self, tensor):
c, w, h = tensor.size()
w1, w2 = self.get_range(w)
h1, h2 = self.get_range(h)
if self.channel_drop_proba == 1:
out = TF.erase(tensor,
h1, w1,
h2-h1, w2-w1,
v=self.default_values,
inplace=self.inplace)
else:
out = tensor if self.inplace else tensor.clone().detach()
for channel, val in enumerate(self.default_values):
if np.random.rand() <= self.channel_drop_proba:
out[channel, w1:w2, h1:h2] = val
return out
def __call__(self, image, label=None):
"""
按照固定大小，随机遮挡图像中的某一块方形区域
:param image: [C,H,W] tensor，必须是tensor
:param label: [H,W] tensor，必须是tensor
:return: [C,H,W] tensor, [H,W] tensor
"""
_, h, w = image.shape
top = random.randint(0, h - self.mask_size) # 随机到遮挡部分的top
left = random.randint(0, w - self.mask_size) # 随机到遮挡部分的left
if random.uniform(0, 1) < 0.5: # 随机遮挡
image = TF.erase(image,
top,
left,
self.mask_size,
self.mask_size,
v=self.value,
inplace=True)
return image, label
def __call__(self, imgs):
"""
Args:
img (Tensor): Tensor image of size (C, H, W) to be erased.
Returns:
img (Tensor): Erased Tensor image.
"""
process_imgs = []
if random.uniform(0, 1) < self.p:
img = imgs[0]
x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=self.value)
for img in imgs:
process_imgs.append(F.erase(img, x, y, h, w, v, self.inplace))
if len(process_imgs) == 0:
return imgs
else:
return process_imgs
return imgs
def __call__(self, imgs):
"""
Args:
imgs (Tensor): Tensor image sequence of size (channel*time*height*width) to be erased.
Returns:
Tensor: Erased Tensor image sequence.
"""
if random.uniform(0, 1) < self.p:
x, y, h, w, v = self.get_params(imgs[:, 0, :, :],
scale=self.scale,
ratio=self.ratio,
value=self.value)
# Apply to all images
output_imgs = []
for I in imgs.transpose(0, 1):
output_imgs.append(
FF.erase(I, x, y, h, w, v, self.inplace).unsqueeze(1))
return torch.cat(output_imgs, dim=1)
return imgs
def random_occlusion(self, img, img_nogt, new_gt):
"""
erase(img: torch.Tensor,
i: int,
j: int,
h: int,
w: int,
v: torch.Tensor,
inplace: bool = False)
"""
occ_i, occ_j, MaxRow, MaxCol = self.affine_img_mask()
tensor_img = TF.pil_to_tensor(img)
tensor_img_nogt = TF.pil_to_tensor(img_nogt)
img_crop_resize = self.crop_resize(tensor_img_nogt, occ_i, occ_j,
MaxRow, MaxCol, new_gt)
img_erase = TF.erase(tensor_img, occ_i, occ_j, MaxRow - occ_i,
MaxCol - occ_j, 0)
self.all_param.append((occ_i, occ_j, MaxRow, MaxCol))
return img_erase, img_crop_resize
def apply_transform(self, img, mask, current_transform=None):
if current_transform is None:
current_transform = self.transform
if isinstance(current_transform, (transforms.Compose)):
for transform in current_transform.transforms:
img, mask = self.apply_transform(img, mask, transform)
elif isinstance(current_transform, (transforms.RandomApply)):
if current_transform.p >= random.random():
img, mask = self.apply_transform(img, mask,
current_transform.transforms)
elif isinstance(current_transform, (transforms.RandomChoice)):
t = random.choice(current_transform.transforms)
img, mask = self.apply_transform(img, mask, t)
elif isinstance(current_transform, (transforms.RandomOrder)):
order = list(range(len(current_transform.transforms)))
random.shuffle(order)
for i in order:
img, mask = self.apply_transform(
img, mask, current_transform.transforms[i])
elif isinstance(
current_transform,
(
transforms.CenterCrop,
transforms.FiveCrop,
transforms.TenCrop,
transforms.ToTensor,
transforms.Grayscale,
transforms.Resize,
),
):
img = current_transform(img)
mask = current_transform(mask)
elif isinstance(
current_transform,
(transforms.Normalize, transforms.Lambda, transforms.Pad)):
img = current_transform(img)
# mask = current_transform(mask) # apply on input only
elif isinstance(current_transform, (transforms.ColorJitter)):
transform = current_transform.get_params(
current_transform.brightness,
current_transform.contrast,
current_transform.saturation,
current_transform.hue,
)
for lambda_transform in transform.transforms:
img = lambda_transform(img)
elif isinstance(current_transform, (transforms.RandomAffine)):
ret = current_transform.get_params(
current_transform.degrees,
current_transform.translate,
current_transform.scale,
current_transform.shear,
img.size,
)
img = F.affine(
img,
*ret,
resample=current_transform.resample,
fillcolor=current_transform.fillcolor,
)
mask = F.affine(
mask,
*ret,
resample=current_transform.resample,
fillcolor=current_transform.fillcolor,
)
elif isinstance(current_transform, (transforms.RandomCrop)):
i, j, h, w = current_transform.get_params(img,
current_transform.size)
img = F.crop(img, i, j, h, w)
mask = F.crop(mask, i, j, h, w)
elif isinstance(current_transform, (transforms.RandomHorizontalFlip)):
if random.random() < current_transform.p:
img = F.hflip(img)
mask = F.hflip(mask)
elif isinstance(current_transform, (transforms.RandomVerticalFlip)):
if random.random() < current_transform.p:
img = F.vflip(img)
mask = F.vflip(mask)
elif isinstance(current_transform, (transforms.RandomPerspective)):
if random.random() < current_transform.p:
width, height = img.size
startpoints, endpoints = current_transform.get_params(
width, height, current_transform.distortion_scale)
img = F.perspective(img, startpoints, endpoints,
current_transform.interpolation)
mask = F.perspective(mask, startpoints, endpoints,
current_transform.interpolation)
elif isinstance(current_transform, (transforms.RandomResizedCrop)):
ret = current_transform.get_params(img, current_transform.scale,
current_transform.ratio)
img = F.resized_crop(img, *ret, current_transform.size,
current_transform.interpolation)
mask = F.resized_crop(mask, *ret, current_transform.size,
current_transform.interpolation)
elif isinstance(current_transform, (transforms.RandomRotation)):
angle = current_transform.get_params(current_transform.degrees)
img = F.rotate(
img,
angle,
current_transform.resample,
current_transform.expand,
current_transform.center,
)
mask = F.rotate(
mask,
angle,
current_transform.resample,
current_transform.expand,
current_transform.center,
)
elif isinstance(current_transform, (transforms.RandomErasing)):
if random.uniform(0, 1) < current_transform.p:
x, y, h, w, v = current_transform.get_params(
img,
scale=current_transform.scale,
ratio=current_transform.ratio,
value=current_transform.value,
)
img = F.erase(img, x, y, h, w, v, current_transform.inplace)
# mask = F.erase(mask, x, y, h, w, v, current_transform.inplace)
else:
raise NotImplementedError(
f'Transform "{current_transform}" not implemented yet')
return img, mask
def __call__(self, img):
import torchvision.transforms.functional as F
img_c, img_h, img_w = img.shape
return F.erase(img, img_h // 2 - self.padding,
img_w // 2 - self.padding, img_h // 2 + self.padding,
img_w // 2 + self.padding, 0, False)