guocheng66's picture
Upload 6 files
57f6383
raw
history blame
5.81 kB
import gradio as gr
from PIL import ImageColor
import onnxruntime
import cv2
import numpy as np
# The common resume photo size is 35mmx45mm
RESUME_PHOTO_W = 350
RESUME_PHOTO_H = 450
# modified from https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/yunet.py
class YuNet:
def __init__(
self,
modelPath,
inputSize=[320, 320],
confThreshold=0.6,
nmsThreshold=0.3,
topK=5000,
backendId=0,
targetId=0,
):
self._modelPath = modelPath
self._inputSize = tuple(inputSize) # [w, h]
self._confThreshold = confThreshold
self._nmsThreshold = nmsThreshold
self._topK = topK
self._backendId = backendId
self._targetId = targetId
self._model = cv2.FaceDetectorYN.create(
model=self._modelPath,
config="",
input_size=self._inputSize,
score_threshold=self._confThreshold,
nms_threshold=self._nmsThreshold,
top_k=self._topK,
backend_id=self._backendId,
target_id=self._targetId,
)
@property
def name(self):
return self.__class__.__name__
def setBackendAndTarget(self, backendId, targetId):
self._backendId = backendId
self._targetId = targetId
self._model = cv2.FaceDetectorYN.create(
model=self._modelPath,
config="",
input_size=self._inputSize,
score_threshold=self._confThreshold,
nms_threshold=self._nmsThreshold,
top_k=self._topK,
backend_id=self._backendId,
target_id=self._targetId,
)
def setInputSize(self, input_size):
self._model.setInputSize(tuple(input_size))
def infer(self, image):
# Forward
faces = self._model.detect(image)
return faces[1]
class ONNXModel:
def __init__(self, model_path, input_w, input_h):
self.model = onnxruntime.InferenceSession(model_path)
self.input_w = input_w
self.input_h = input_h
def preprocess(self, rgb, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
# convert the input data into the float32 input
img_data = (
np.array(cv2.resize(rgb, (self.input_w, self.input_h)))
.transpose(2, 0, 1)
.astype("float32")
)
# normalize
norm_img_data = np.zeros(img_data.shape).astype("float32")
for i in range(img_data.shape[0]):
norm_img_data[i, :, :] = img_data[i, :, :] / 255
norm_img_data[i, :, :] = (norm_img_data[i, :, :] - mean[i]) / std[i]
# add batch channel
norm_img_data = norm_img_data.reshape(1, 3, self.input_h, self.input_w).astype(
"float32"
)
return norm_img_data
def forward(self, image):
input_data = self.preprocess(image)
output_data = self.model.run(["argmax_0.tmp_0"], {"x": input_data})
return output_data
def make_resume_photo(rgb, background_color):
h, w, _ = rgb.shape
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
# Initialize models
face_detector = YuNet("models/face_detection_yunet_2023mar.onnx")
face_detector.setInputSize([w, h])
human_segmentor = ONNXModel(
"models/human_pp_humansegv2_lite_192x192_inference_model.onnx", 192, 192
)
# yunet uses opencv bgr image format
detections = face_detector.infer(bgr)
results = []
for idx, det in enumerate(detections):
# bounding box
pt1 = np.array((det[0], det[1]))
pt2 = np.array((det[0] + det[2], det[1] + det[3]))
# face landmarks
landmarks = det[4:14].reshape((5, 2))
right_eye = landmarks[0]
left_eye = landmarks[1]
angle = np.arctan2(right_eye[1] - left_eye[1], (right_eye[0] - left_eye[0]))
rmat = cv2.getRotationMatrix2D((0, 0), -angle, 1)
# apply rotation
rotated_bgr = cv2.warpAffine(bgr, rmat, (bgr.shape[1], bgr.shape[0]))
rotated_pt1 = rmat[:, :-1] @ pt1
rotated_pt2 = rmat[:, :-1] @ pt2
face_w, face_h = rotated_pt2 - rotated_pt1
up_length = int(face_h / 4)
down_length = int(face_h / 3)
crop_h = face_h + up_length + down_length
crop_w = int(crop_h * (RESUME_PHOTO_W / RESUME_PHOTO_H))
pt1 = np.array(
(rotated_pt1[0] - (crop_w - face_w) / 2, rotated_pt1[1] - up_length)
).astype(np.int32)
pt2 = np.array((pt1[0] + crop_w, pt1[1] + crop_h)).astype(np.int32)
resume_photo = rotated_bgr[pt1[1] : pt2[1], pt1[0] : pt2[0], :]
rgb = cv2.cvtColor(resume_photo, cv2.COLOR_BGR2RGB)
mask = human_segmentor.forward(rgb)
mask = mask[0].transpose(1, 2, 0)
mask = cv2.resize(
mask.astype(np.uint8), (resume_photo.shape[1], resume_photo.shape[0])
)
resume_photo = cv2.cvtColor(resume_photo, cv2.COLOR_BGR2RGB)
resume_photo[mask == 0] = ImageColor.getcolor(background_color, "RGB")
resume_photo = cv2.resize(resume_photo, (RESUME_PHOTO_W, RESUME_PHOTO_H))
results.append(resume_photo)
return results
title = "Resume Photo Maker"
demo = gr.Interface(
fn=make_resume_photo,
inputs=[
gr.Image(type="numpy", label="input"),
gr.ColorPicker(label="background color"),
],
outputs=gr.Gallery(label="output"),
examples=[
["images/elon.jpg", "#FFFFFF"],
["images/9_Press_Conference_Press_Conference_9_45.jpg", "#FFFFFF"],
],
title=title,
allow_flagging="never",
article="<p style='text-align: center;'><a href='https://github.com/bot66/resume-photo-maker' target='_blank'>Github Repo</a></p>",
)
if __name__ == "__main__":
demo.launch()