|
from torchvision.utils import make_grid |
|
from PIL import Image, ImageDraw, ImageFont |
|
import numpy as np |
|
import torch |
|
import math |
|
import cv2 |
|
def make_grid_(imgs, save_file, nrow=10, pad_value=1): |
|
if isinstance(imgs, list): |
|
if isinstance(imgs[0], Image.Image): |
|
imgs = [torch.from_numpy(np.array(img)/255.) for img in imgs] |
|
elif isinstance(imgs[0], np.ndarray): |
|
imgs = [torch.from_numpy(img/255.) for img in imgs] |
|
imgs = torch.stack(imgs, 0).permute(0, 3, 1, 2) |
|
if isinstance(imgs, np.ndarray): |
|
imgs = torch.from_numpy(imgs) |
|
|
|
img_grid = make_grid(imgs, nrow=nrow, padding=2, pad_value=pad_value) |
|
img_grid = img_grid.permute(1, 2, 0).numpy() |
|
img_grid = (img_grid * 255).astype(np.uint8) |
|
img_grid = Image.fromarray(img_grid) |
|
img_grid.save(save_file) |
|
|
|
def draw_caption(img, text, pos, size=100, color=(128, 128, 128)): |
|
draw = ImageDraw.Draw(img) |
|
|
|
font = ImageFont.load_default() |
|
font = font.font_variant(size=size) |
|
draw.text(pos, text, color, font=font) |
|
return img |
|
|
|
|
|
def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]): |
|
|
|
stickwidth = 4 |
|
limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) |
|
kps = np.array(kps) |
|
|
|
w, h = image_pil.size |
|
out_img = np.zeros([h, w, 3]) |
|
|
|
for i in range(len(limbSeq)): |
|
index = limbSeq[i] |
|
color = color_list[index[0]] |
|
|
|
x = kps[index][:, 0] |
|
y = kps[index][:, 1] |
|
length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 |
|
angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) |
|
polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1) |
|
out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) |
|
out_img = (out_img * 0.6).astype(np.uint8) |
|
|
|
for idx_kp, kp in enumerate(kps): |
|
color = color_list[idx_kp] |
|
x, y = kp |
|
out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) |
|
|
|
out_img_pil = Image.fromarray(out_img.astype(np.uint8)) |
|
return out_img_pil |
|
|
|
def resize_img(input_image, max_side=1280, min_side=1024, size=None, |
|
pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64): |
|
|
|
w, h = input_image.size |
|
if size is not None: |
|
w_resize_new, h_resize_new = size |
|
else: |
|
ratio = min_side / min(h, w) |
|
w, h = round(ratio*w), round(ratio*h) |
|
ratio = max_side / max(h, w) |
|
input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode) |
|
w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number |
|
h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number |
|
input_image = input_image.resize([w_resize_new, h_resize_new], mode) |
|
|
|
if pad_to_max_side: |
|
res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255 |
|
offset_x = (max_side - w_resize_new) // 2 |
|
offset_y = (max_side - h_resize_new) // 2 |
|
res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image) |
|
input_image = Image.fromarray(res) |
|
return input_image |
