File size: 2,962 Bytes
c0aacde 0fd2553 c0aacde 0fd2553 c0aacde 0fd2553 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import cv2
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from math import atan2
from PIL import Image as PImage
from ultralytics import YOLO
OUT_W = 130
OUT_H = 170
OUT_EYE_SPACE = 64
OUT_FACE_WIDTH = 89
OUT_NOSE_TOP = 72
EYE_0_IDX = 36
EYE_1_IDX = 45
TEMPLE_0_IDX = 0
TEMPLE_1_IDX = 16
yolo_model_path = hf_hub_download(repo_id="AdamCodd/YOLOv11n-face-detection", filename="model.pt")
face_detector = YOLO(yolo_model_path)
LBFmodel = "./models/lbfmodel.yaml"
landmark_detector = cv2.face.createFacemarkLBF()
landmark_detector.loadModel(LBFmodel)
NUM_OUTS = 16
all_outputs = [gr.Image(format="jpeg", visible=False) for _ in range(NUM_OUTS)]
def face(img_in):
out_pad = NUM_OUTS * [gr.Image(visible=False)]
if img_in is None:
return out_pad
img = img_in.copy()
img.thumbnail((1000,1000))
img_np = np.array(img).copy()
iw,ih = img.size
output = face_detector.predict(img, verbose=False)
if len(output) < 1 or len(output[0]) < 1:
return out_pad
faces_xyxy = output[0].boxes.xyxy.numpy()
faces = np.array([[x0, y0, (x1 - x0), (y1 - y0)] for x0,y0,x1,y1 in faces_xyxy])
biggest_faces = faces[np.argsort(-faces[:,2])]
_, landmarks = landmark_detector.fit(img_np, biggest_faces)
if len(landmarks) < 1:
return out_pad
out_images = []
for landmark in landmarks:
eye0 = np.array(landmark[0][EYE_0_IDX])
eye1 = np.array(landmark[0][EYE_1_IDX])
temple0 = np.array(landmark[0][TEMPLE_0_IDX])
temple1 = np.array(landmark[0][TEMPLE_1_IDX])
mid = np.mean([eye0, eye1], axis=0)
eye_line = eye1 - eye0
tilt = atan2(eye_line[1], eye_line[0])
tilt_deg = 180 * tilt / np.pi
scale = min(OUT_EYE_SPACE / np.linalg.norm(eye1 - eye0),
OUT_FACE_WIDTH / np.linalg.norm(temple1 - temple0))
img_s = img.resize((int(iw * scale), int(ih * scale)))
# rotate around nose
new_mid = [int(c * scale) for c in mid]
crop_box = (new_mid[0] - (OUT_W // 2),
new_mid[1] - OUT_NOSE_TOP,
new_mid[0] + (OUT_W // 2),
new_mid[1] + (OUT_H - OUT_NOSE_TOP))
img_out = img_s.rotate(tilt_deg, center=new_mid, resample=PImage.Resampling.BICUBIC).crop(crop_box).convert("L")
out_images.append(gr.Image(img_out, visible=True))
out_images += out_pad
return out_images[:NUM_OUTS]
with gr.Blocks() as demo:
gr.Markdown("""
# PSAM 5020 Face Alignment Tool.
## Interface for face detection, alignment, cropping\
to help create dataset for [WK12](https://github.com/PSAM-5020-2025F-A/WK12) / [HW12](https://github.com/PSAM-5020-2025F-A/Homework12).
""")
gr.Interface(
face,
inputs=gr.Image(type="pil"),
outputs=all_outputs,
flagging_mode="never",
cache_examples=True,
examples=[["./imgs/03.webp"], ["./imgs/11.jpg"], ["./imgs/people.jpg"]],
)
if __name__ == "__main__":
demo.launch(ssr_mode=False)
|