230 lines
7.3 KiB
Python
230 lines
7.3 KiB
Python
import math
|
|
import torch
|
|
import torch as th
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from scipy.interpolate import interp1d
|
|
import importlib
|
|
from matplotlib.path import Path
|
|
|
|
|
|
def get_polygon_mask_params(
|
|
mask_size, box, num_vertices, mask_scale, min_scale, max_scale
|
|
):
|
|
center = ((box[2] + box[0]) / 2, (box[3] + box[1]) / 2)
|
|
sizes = (box[2] - box[0], box[3] - box[1])
|
|
|
|
part_avg_radii = np.linspace(
|
|
mask_scale * sizes[0] / 2, mask_scale * sizes[1] / 2, num_vertices // 4
|
|
)
|
|
part_avg_radii = np.clip(
|
|
part_avg_radii, min_scale * min(mask_size), max_scale * min(mask_size)
|
|
)
|
|
avg_radii = np.concatenate(
|
|
[
|
|
part_avg_radii,
|
|
part_avg_radii[::-1],
|
|
part_avg_radii,
|
|
part_avg_radii[::-1],
|
|
]
|
|
)
|
|
return center, avg_radii
|
|
|
|
|
|
def smooth_cerv(x, y):
|
|
num_vertices = x.shape[0]
|
|
x = np.concatenate((x[-3:-1], x, x[1:3]))
|
|
y = np.concatenate((y[-3:-1], y, y[1:3]))
|
|
t = np.arange(x.shape[0])
|
|
|
|
ti = np.linspace(2, num_vertices + 1, 4 * num_vertices)
|
|
xi = interp1d(t, x, kind="quadratic")(ti)
|
|
yi = interp1d(t, y, kind="quadratic")(ti)
|
|
return xi, yi
|
|
|
|
|
|
def get_polygon_mask(mask_size, mask_points):
|
|
x, y = np.meshgrid(np.arange(mask_size[0]), np.arange(mask_size[1]))
|
|
x, y = x.flatten(), y.flatten()
|
|
points = np.vstack((x, y)).T
|
|
|
|
path = Path(mask_points)
|
|
grid = path.contains_points(points)
|
|
grid = grid.reshape((mask_size[0], mask_size[1]))
|
|
return 1.0 - grid.astype(np.int32)
|
|
|
|
|
|
def generate_polygon(
|
|
mask_size, center, num_vertices, radii, radii_var, angle_var, smooth=True
|
|
):
|
|
angle_steps = np.random.uniform(
|
|
1.0 - angle_var, 1.0 + angle_var, size=(num_vertices,)
|
|
)
|
|
angle_steps = 2 * np.pi * angle_steps / angle_steps.sum()
|
|
|
|
radii = np.random.normal(radii, radii_var * radii)
|
|
radii = np.clip(radii, 0, 2 * radii)
|
|
angles = np.cumsum(angle_steps)
|
|
x = center[0] + radii * np.cos(angles)
|
|
y = center[1] + radii * np.sin(angles)
|
|
|
|
if smooth:
|
|
x, y = smooth_cerv(x, y)
|
|
points = np.concatenate((x.reshape(-1, 1), y.reshape(-1, 1)), axis=-1)
|
|
points = list(map(tuple, points.tolist()))
|
|
return get_polygon_mask(mask_size, points)
|
|
|
|
|
|
def generate_circle_frame(mask_size, side_scales, num_vertices, radii_var, smooth=True):
|
|
num_vertices_per_side = num_vertices // 4
|
|
x_size, y_size = mask_size
|
|
up_radii = np.array([y_size * (1.0 - side_scales[0]) // 2] * num_vertices_per_side)
|
|
down_radii = np.array(
|
|
[y_size * (1.0 - side_scales[1]) // 2] * num_vertices_per_side
|
|
)
|
|
left_radii = np.array(
|
|
[x_size * (1.0 - side_scales[2]) // 2] * num_vertices_per_side
|
|
)
|
|
right_radii = np.array(
|
|
[x_size * (1.0 - side_scales[3]) // 2] * num_vertices_per_side
|
|
)
|
|
|
|
center = (x_size // 2, y_size // 2)
|
|
radii = np.concatenate(
|
|
[
|
|
right_radii[num_vertices_per_side // 2 :],
|
|
down_radii,
|
|
left_radii,
|
|
up_radii,
|
|
right_radii[: num_vertices_per_side // 2],
|
|
]
|
|
)
|
|
return 1.0 - generate_polygon(
|
|
mask_size, center, num_vertices, radii, radii_var, 0.0, smooth=smooth
|
|
)
|
|
|
|
|
|
def generate_square_frame(mask_size, side_scales, num_vertices, radii_var, smooth=True):
|
|
num_vertices_per_side = num_vertices // 4
|
|
x_size, y_size = mask_size
|
|
diag_size = np.sqrt(x_size**2 + y_size**2)
|
|
|
|
up_radii = np.linspace(
|
|
diag_size * (1.0 - side_scales[0]) // 2,
|
|
y_size * (1.0 - side_scales[0]) // 2,
|
|
num_vertices_per_side // 2,
|
|
)
|
|
down_radii = np.linspace(
|
|
diag_size * (1.0 - side_scales[1]) // 2,
|
|
y_size * (1.0 - side_scales[1]) // 2,
|
|
num_vertices_per_side // 2,
|
|
)
|
|
left_radii = np.linspace(
|
|
diag_size * (1.0 - side_scales[2]) // 2,
|
|
x_size * (1.0 - side_scales[2]) // 2,
|
|
num_vertices_per_side // 2,
|
|
)
|
|
right_radii = np.linspace(
|
|
diag_size * (1.0 - side_scales[3]) // 2,
|
|
x_size * (1.0 - side_scales[3]) // 2,
|
|
num_vertices_per_side // 2,
|
|
)
|
|
|
|
center = (x_size // 2, y_size // 2)
|
|
radii = np.concatenate(
|
|
[
|
|
right_radii[::-1],
|
|
down_radii,
|
|
down_radii[::-1],
|
|
left_radii,
|
|
left_radii[::-1],
|
|
up_radii,
|
|
up_radii[::-1],
|
|
right_radii,
|
|
]
|
|
)
|
|
return 1.0 - generate_polygon(
|
|
mask_size, center, num_vertices, radii, radii_var, 0.0, smooth=smooth
|
|
)
|
|
|
|
|
|
def generate_mask(mask_size, box, box_prob=0.1):
|
|
mask = np.ones(mask_size)
|
|
if np.random.binomial(1, box_prob):
|
|
box = [int(i) for i in box]
|
|
mask[box[1] : box[3], box[0] : box[2]] = 0
|
|
|
|
else:
|
|
actions = np.random.randint(0, 2, (2,))
|
|
if 0 in actions:
|
|
num_vertices = 16
|
|
center, radii = get_polygon_mask_params(
|
|
mask_size,
|
|
box,
|
|
num_vertices,
|
|
mask_scale=1.5,
|
|
min_scale=0.1,
|
|
max_scale=0.6,
|
|
)
|
|
mask *= generate_polygon(
|
|
mask_size, center, num_vertices, radii, radii_var=0.15, angle_var=0.15
|
|
)
|
|
if 1 in actions:
|
|
radii_var = 0.15 * np.random.random()
|
|
num_vertices = np.random.choice([16, 32])
|
|
if np.random.random() < 0.5:
|
|
side_scales = 0.25 * np.random.random((4,)) + 0.05
|
|
mask *= generate_square_frame(
|
|
mask_size, side_scales, num_vertices, radii_var=radii_var
|
|
)
|
|
else:
|
|
side_scales = 0.15 * np.random.random((4,)) + 0.1
|
|
mask *= generate_circle_frame(
|
|
mask_size, side_scales, num_vertices, radii_var=radii_var
|
|
)
|
|
return mask
|
|
|
|
|
|
def get_boxes(bs, target_size, min_scale=0.1, max_scale=0.62):
|
|
min_size_x = min_scale * target_size[0]
|
|
max_size_x = max_scale * target_size[0]
|
|
min_size_y = min_scale * target_size[1]
|
|
max_size_y = max_scale * target_size[1]
|
|
|
|
boxes_size_x = (max_size_x - min_size_x) * np.random.random((bs, 1)) + min_size_x
|
|
boxes_size_y = (max_size_y - min_size_y) * np.random.random((bs, 1)) + min_size_y
|
|
|
|
x0 = (target_size[0] - max_size_x) * np.random.random((bs, 1))
|
|
y0 = (target_size[1] - max_size_y) * np.random.random((bs, 1))
|
|
|
|
boxes = np.concatenate((x0, y0, x0 + boxes_size_x, y0 + boxes_size_x), -1)
|
|
return boxes.tolist()
|
|
|
|
|
|
def get_image_mask(bs, target_size):
|
|
boxes = get_boxes(bs, target_size)
|
|
image_mask = torch.stack(
|
|
[torch.tensor(generate_mask(target_size, box)) for box in boxes]
|
|
)
|
|
return image_mask
|
|
|
|
|
|
def freeze_decoder(
|
|
model,
|
|
freeze_resblocks=False,
|
|
freeze_attention=False,
|
|
):
|
|
for name, p in model.named_parameters():
|
|
name = name.lower()
|
|
if (
|
|
"in_layers" in name
|
|
or "h_upd" in name
|
|
or "x_upd" in name
|
|
or "emb_layers" in name
|
|
or "out_layers" in name
|
|
):
|
|
p.requires_grad = not freeze_resblocks
|
|
elif "proj_out" in name or "qkv" in name:
|
|
p.requires_grad = not freeze_attention
|
|
return model
|