File size: 4,369 Bytes
ec0fdfd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
import cv2
from PIL import Image
import random
###################################################################
# random mask generation
###################################################################
def random_regular_mask(img):
"""Generate a random regular mask
:param img: original image size C*H*W
:return: mask
"""
mask = torch.ones_like(img)[0:1, :, :]
s = img.size()
N_mask = random.randint(1, 5)
lim_x = s[1] - s[1] / (N_mask + 1)
lim_y = s[2] - s[2] / (N_mask + 1)
for _ in range(N_mask):
x = random.randint(0, int(lim_x))
y = random.randint(0, int(lim_y))
range_x = x + random.randint(int(s[1] / (N_mask + 7)), min(int(s[1] - x), int(s[1] / 2)))
range_y = y + random.randint(int(s[2] / (N_mask + 7)), min(int(s[2] - y), int(s[2] / 2)))
mask[:, int(x) : int(range_x), int(y) : int(range_y)] = 0
return mask
def center_mask(img):
"""Generate a center hole with 1/4*W and 1/4*H
:param img: original image size C*H*W
:return: mask
"""
mask = torch.ones_like(img)[0:1, :, :]
s = img.size()
mask[:, int(s[1]/4):int(s[1]*3/4), int(s[2]/4):int(s[2]*3/4)] = 0
return mask
def random_irregular_mask(img):
"""Generate a random irregular mask with lines, circles and ellipses
:param img: original image size C*H*W
:return: mask
"""
transform = transforms.Compose([transforms.ToTensor()])
mask = torch.ones_like(img)[0:1, :, :]
s = mask.size()
img = np.zeros((s[1], s[2], 1), np.uint8)
max_width = int(min(s[1]/10, s[2]/10))
N_mask = random.randint(16, 64)
for _ in range(N_mask):
model = random.random()
if model < 0.2: # Draw random lines
x1, x2 = random.randint(1, s[1]), random.randint(1, s[1])
y1, y2 = random.randint(1, s[2]), random.randint(1, s[2])
thickness = random.randint(2, max_width)
cv2.line(img, (x1, y1), (x2, y2), (1, 1, 1), thickness)
elif (model > 0.2 and model < 0.5): # Draw random circles
x1, y1 = random.randint(1, s[1]), random.randint(1, s[2])
radius = random.randint(2, max_width)
cv2.circle(img, (x1, y1), radius, (1, 1, 1), -1)
else: # draw random ellipses
x1, y1 = random.randint(1, s[1]), random.randint(1, s[2])
s1, s2 = random.randint(1, s[1]), random.randint(1, s[2])
a1, a2, a3 = random.randint(3, 180), random.randint(3, 180), random.randint(3, 180)
thickness = random.randint(2, max_width)
cv2.ellipse(img, (x1, y1), (s1, s2), a1, a2, a3, (1, 1, 1), thickness)
img = img.reshape(s[2], s[1])
img = Image.fromarray(img*255)
img_mask = transform(img)
for j in range(s[0]):
mask[j, :, :] = img_mask
return mask
def scale_img(img, size):
h_ratio = img.size(-1) // size[-1]
w_ratio = img.size(-2) // size[-2]
scaled_img = F.avg_pool2d(img, kernel_size=(w_ratio, h_ratio), stride=(w_ratio, h_ratio))
return scaled_img
def scale_pyramid(img, num_scales):
scaled_imgs = [img]
for i in range(1, num_scales):
ratio = 2**i
scaled_img = F.avg_pool2d(img, kernel_size=ratio, stride=ratio)
scaled_imgs.append(scaled_img)
scaled_imgs.reverse()
return scaled_imgs
def jacobian(y, x, point=None, create_graph=True):
"""Calculate the jacobian matrix for given point"""
jac = []
flat_y = y.reshape(-1)
b, c, h, w = y.size()
if point is not None:
i = point[0] * h + point[1]
input_y = flat_y[i]
grad_x = torch.autograd.grad(input_y, x, retain_graph=True, grad_outputs=torch.ones(input_y.size()).to(x.device),
create_graph=create_graph, only_inputs=True)[0]
jac.append(grad_x.reshape(x.shape))
return jac
else:
for i in range(len(flat_y)):
input_y = flat_y[i]
grad_x = torch.autograd.grad(input_y, x, retain_graph=True, grad_outputs=torch.ones(input_y.size()).to(x.device),
create_graph=create_graph, only_inputs=True)[0]
jac.append(grad_x.reshape(x.shape))
return torch.stack(jac).reshape(y.shape + x.shape) |