Spaces:
Sleeping
Sleeping
from transformers import DetrImageProcessor, DetrForObjectDetection | |
from PIL import Image, ImageDraw | |
import torch | |
import gradio as gr | |
import requests | |
from io import BytesIO | |
# Load pre-trained DETR model | |
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") | |
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") | |
# COCO class index for "person" = 1 (used as proxy for face detection) | |
FACE_CLASS_INDEX = 1 | |
def detect_faces(img: Image.Image): | |
# Prepare input for the model | |
inputs = processor(images=img, return_tensors="pt") | |
outputs = model(**inputs) | |
# Get outputs | |
target_sizes = torch.tensor([img.size[::-1]]) | |
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0] | |
# Draw bounding boxes | |
draw = ImageDraw.Draw(img) | |
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): | |
if label.item() == FACE_CLASS_INDEX: # 'person' | |
box = [round(i, 2) for i in box.tolist()] | |
draw.rectangle(box, outline="green", width=3) | |
draw.text((box[0], box[1]), f"{score:.2f}", fill="green") | |
return img | |
# Gradio interface | |
iface = gr.Interface( | |
fn=detect_faces, | |
inputs=gr.Image(type="pil"), | |
outputs="image", | |
title="Face Detection App (Hugging Face + Gradio)", | |
description="Upload an image and detect faces using facebook/detr-resnet-50 model." | |
) | |
iface.launch() | |