139 lines
3.7 KiB
Python
139 lines
3.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import cv2
|
|
import numpy as np
|
|
import torchvision.transforms.functional as TF
|
|
|
|
|
|
def get_bbox_from_mask(mask):
|
|
h, w = mask.shape[0], mask.shape[1]
|
|
if mask.sum() < 10:
|
|
return 0, h, 0, w
|
|
rows = np.any(mask, axis=1)
|
|
cols = np.any(mask, axis=0)
|
|
y1, y2 = np.where(rows)[0][[0, -1]]
|
|
x1, x2 = np.where(cols)[0][[0, -1]]
|
|
return (y1, y2, x1, x2)
|
|
|
|
|
|
def pad_to_square(image, pad_value=255, random=False):
|
|
H, W = image.shape[0], image.shape[1]
|
|
if H == W:
|
|
return image, 0, 0
|
|
|
|
padd = abs(H - W)
|
|
if random:
|
|
padd_1 = int(np.random.randint(0, padd))
|
|
else:
|
|
padd_1 = int(padd / 2)
|
|
padd_2 = padd - padd_1
|
|
|
|
if H > W:
|
|
pad_param = ((0, 0), (padd_1, padd_2), (0, 0))
|
|
else:
|
|
pad_param = ((padd_1, padd_2), (0, 0), (0, 0))
|
|
|
|
# print(pad_param, pad_value)
|
|
image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
|
|
return image, padd_1, padd_2
|
|
|
|
|
|
def box_in_box(small_box, big_box):
|
|
y1, y2, x1, x2 = small_box
|
|
y1_b, _, x1_b, _ = big_box
|
|
y1, y2, x1, x2 = y1 - y1_b, y2 - y1_b, x1 - x1_b, x2 - x1_b
|
|
return (y1, y2, x1, x2)
|
|
|
|
|
|
def box2squre(image, box):
|
|
H, W = image.shape[0], image.shape[1]
|
|
y1, y2, x1, x2 = box
|
|
cx = (x1 + x2) // 2
|
|
cy = (y1 + y2) // 2
|
|
h, w = y2 - y1, x2 - x1
|
|
|
|
if h >= w:
|
|
x1 = cx - h // 2
|
|
x2 = x1 + h
|
|
else:
|
|
y1 = cy - w // 2
|
|
y2 = y1 + w
|
|
x1 = max(0, x1)
|
|
x2 = min(W, x2)
|
|
y1 = max(0, y1)
|
|
y2 = min(H, y2)
|
|
return (y1, y2, x1, x2)
|
|
|
|
|
|
def expand_bbox(mask,
|
|
yyxx,
|
|
ratio=1.0,
|
|
min_crop=0,
|
|
expand_type='center',
|
|
to_square=False):
|
|
y1, y2, x1, x2 = yyxx
|
|
h = y2 - y1 + 1
|
|
w = x2 - x1 + 1
|
|
|
|
H, W = mask.shape[0], mask.shape[1]
|
|
xc, yc = 0.5 * (x1 + x2), 0.5 * (y1 + y2)
|
|
|
|
def expand(k):
|
|
if isinstance(ratio, tuple) or isinstance(ratio, list):
|
|
r = np.random.uniform(*ratio)
|
|
k = k * r
|
|
else:
|
|
k = ratio * k
|
|
return k
|
|
|
|
new_h = expand(h)
|
|
new_w = expand(w)
|
|
new_h = max(new_h, min_crop)
|
|
new_w = max(new_w, min_crop)
|
|
|
|
if to_square:
|
|
if new_w / new_h < 0.334:
|
|
new_w = new_w + 1.0 / 3.0 * new_h
|
|
elif new_h / new_w < 0.334:
|
|
new_h = new_h + 1.0 / 3.0 * new_w
|
|
|
|
if expand_type == 'center':
|
|
x1 = max(0, int(xc - new_w * 0.5))
|
|
x2 = min(W, int(xc + new_w * 0.5))
|
|
y1 = max(0, int(yc - new_h * 0.5))
|
|
y2 = min(H, int(yc + new_h * 0.5))
|
|
else:
|
|
x1 = max(0, min(x1,
|
|
int(x2 - new_w * np.random.uniform(w / new_w, 1.0))))
|
|
x2 = min(W, max(x2, x1 + new_w))
|
|
y1 = max(0, min(y1,
|
|
int(y2 - new_h * np.random.uniform(h / new_h, 1.0))))
|
|
y2 = min(H, max(y2, y1 + new_h))
|
|
|
|
return (int(y1), int(y2), int(x1), int(x2))
|
|
|
|
|
|
def crop_back(pred, tar_image, extra_sizes, tar_box_yyxx_crop):
|
|
H1, W1, H2, W2, pad1, pad2 = extra_sizes
|
|
y1, y2, x1, x2 = tar_box_yyxx_crop
|
|
pred = TF.resize(pred, (H2, W2), antialias=True)
|
|
|
|
if W1 < W2:
|
|
# pad width
|
|
assert H1 == H2 and (pad1 + W1) == (W2 - pad2)
|
|
pred = pred[:, :, pad1 + 2:(W2 - pad2 - 2)]
|
|
tar_image[:, y1:y2, x1 + 2:x2 - 2] = pred
|
|
elif H1 < H2:
|
|
# pad height
|
|
assert W1 == W2 and (pad1 + H1) == (H2 - pad2)
|
|
pred = pred[:, pad1 + 2:(H2 - pad2 - 2), :]
|
|
tar_image[:, y1 + 2:y2 - 2, x1:x2] = pred
|
|
else:
|
|
tar_image[:, y1:y2, x1:x2] = pred
|
|
return tar_image
|
|
|
|
|
|
def save_image(image, save_path):
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite(save_path, image)
|