Spaces:
Sleeping
Sleeping
File size: 2,663 Bytes
d06fdab | 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 | import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np
# Download and load model
model_path = hf_hub_download(repo_id="chuuhtetnaing/face-detection-yolo26n", filename="best.pt")
model = YOLO(model_path)
def crop_face_with_padding(image, bbox, padding_ratio=0.15):
"""Crop face with padding, handling edge cases."""
img_w, img_h = image.size
x1, y1, x2, y2 = bbox
box_w = x2 - x1
box_h = y2 - y1
pad_w = int(box_w * padding_ratio)
pad_h = int(box_h * padding_ratio)
# Apply padding with boundary checks
new_x1 = max(0, int(x1 - pad_w))
new_y1 = max(0, int(y1 - pad_h))
new_x2 = min(img_w, int(x2 + pad_w))
new_y2 = min(img_h, int(y2 + pad_h))
return image.crop((new_x1, new_y1, new_x2, new_y2))
def detect_faces(image):
"""Detect faces and return annotated image + cropped faces."""
if image is None:
return None, None
image = Image.fromarray(image) if isinstance(image, np.ndarray) else image
# Run detection
results = model(image)[0]
# Draw bounding boxes
annotated = image.copy()
draw = ImageDraw.Draw(annotated)
cropped_faces = []
for box in results.boxes:
x1, y1, x2, y2 = box.xyxy[0].tolist()
conf = box.conf[0].item()
# Draw rectangle and label
draw.rectangle([x1, y1, x2, y2], outline="green", width=20)
font = ImageFont.load_default(size=200)
draw.text((x1, y1 - 200), f"{conf:.2f}", fill="green", font=font)
# Crop with padding
cropped = crop_face_with_padding(image, (x1, y1, x2, y2))
cropped_faces.append(cropped)
# Create gallery of cropped faces
cropped_gallery = cropped_faces if cropped_faces else None
return annotated, cropped_gallery
with gr.Blocks(css=".main-row { margin-top: 30px !important; }") as demo:
gr.Markdown("<h1 style='text-align: center; margin-top: 20px'>Face Detection with YOLO26n</h1>")
with gr.Row(elem_classes="main-row"):
with gr.Column(scale=1):
input_image = gr.Image(label="Upload Image", type="pil", height=400)
detect_btn = gr.Button("Detect Faces", variant="primary")
with gr.Column(scale=1):
with gr.Row():
output_image = gr.Image(label="Detected Faces", height=400)
cropped_gallery = gr.Gallery(label="Cropped Faces", columns=2, height=400)
detect_btn.click(
fn=detect_faces,
inputs=input_image,
outputs=[output_image, cropped_gallery]
)
if __name__ == "__main__":
demo.launch()
|