94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import torch
|
|
import numpy as np
|
|
from insightface.utils import face_align
|
|
from PIL import Image
|
|
import math
|
|
import cv2
|
|
|
|
def extract_arcface_bgr_embedding(in_image, landmark, arcface_model, in_settings=None):
|
|
kps = landmark
|
|
arc_face_image = face_align.norm_crop(in_image, landmark=np.array(kps), image_size=112)
|
|
arc_face_image = torch.from_numpy(arc_face_image).unsqueeze(0).permute(0,3,1,2) / 255.
|
|
arc_face_image = 2 * arc_face_image - 1
|
|
arc_face_image = arc_face_image.cuda().contiguous()
|
|
face_emb = arcface_model(arc_face_image)[0] # [512], normalized
|
|
return face_emb
|
|
|
|
def tensor_to_np_image(tensor):
|
|
return tensor.mul(255).clamp(0, 255).byte().cpu().numpy()
|
|
|
|
def np_image_to_tensor(image):
|
|
return torch.clamp(torch.from_numpy(image).float() / 255., 0, 1)
|
|
|
|
def resize_and_pad_pil_image(source_img, target_img_size):
|
|
# Get original and target sizes
|
|
source_img_size = source_img.size
|
|
target_width, target_height = target_img_size
|
|
|
|
# Determine the new size based on the shorter side of target_img
|
|
if target_width <= target_height:
|
|
new_width = target_width
|
|
new_height = int(target_width * (source_img_size[1] / source_img_size[0]))
|
|
else:
|
|
new_height = target_height
|
|
new_width = int(target_height * (source_img_size[0] / source_img_size[1]))
|
|
|
|
# Resize the source image using LANCZOS interpolation for high quality
|
|
resized_source_img = source_img.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
# Compute padding to center resized image
|
|
pad_left = (target_width - new_width) // 2
|
|
pad_top = (target_height - new_height) // 2
|
|
|
|
# Create a new image with white background
|
|
padded_img = Image.new("RGB", target_img_size, (255, 255, 255))
|
|
padded_img.paste(resized_source_img, (pad_left, pad_top))
|
|
|
|
return padded_img
|
|
|
|
# modified from https://github.com/instantX-research/InstantID/blob/main/pipeline_stable_diffusion_xl_instantid.py
|
|
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 escape_path_for_url(path):
|
|
return path.replace("\\", "/")
|