pratikshahp's picture
Update app.py
f9d3b76 verified
raw
history blame
1.46 kB
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()