OMNIVER / utils.py
okaris
Omni-Zero
1d4b9ba
raw
history blame
3.27 kB
import math
import PIL
import cv2
import numpy as np
from diffusers.utils import load_image
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])
if type(image_pil) == PIL.Image.Image:
out_img = np.array(image_pil)
else:
out_img = image_pil
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 = PIL.Image.fromarray(out_img.astype(np.uint8))
return out_img_pil
def load_and_resize_image(image_path, max_width, max_height, maintain_aspect_ratio=True):
# Open the image
# image = Image.open(image_path)
image = load_image(image_path)
# Get the current width and height of the image
current_width, current_height = image.size
if maintain_aspect_ratio:
# Calculate the aspect ratio of the image
aspect_ratio = current_width / current_height
# Calculate the new dimensions based on the max width and height
if current_width / max_width > current_height / max_height:
new_width = max_width
new_height = int(new_width / aspect_ratio)
else:
new_height = max_height
new_width = int(new_height * aspect_ratio)
else:
# Use the max width and height as the new dimensions
new_width = max_width
new_height = max_height
# Ensure the new dimensions are divisible by 8
new_width = (new_width // 8) * 8
new_height = (new_height // 8) * 8
# Resize the image
resized_image = image.resize((new_width, new_height))
return resized_image
from PIL import Image
def align_images(image1, image2):
"""
Resize two images to the same dimensions by cropping the larger image(s) to match the smaller one.
Args:
image1 (PIL.Image): First image to be aligned.
image2 (PIL.Image): Second image to be aligned.
Returns:
tuple: A tuple containing two images with the same dimensions.
"""
# Determine the new size by taking the smaller width and height from both images
new_width = min(image1.size[0], image2.size[0])
new_height = min(image1.size[1], image2.size[1])
# Crop both images if necessary
if image1.size != (new_width, new_height):
image1 = image1.crop((0, 0, new_width, new_height))
if image2.size != (new_width, new_height):
image2 = image2.crop((0, 0, new_width, new_height))
return image1, image2