Spaces:
Sleeping
Sleeping
from transformers import DetrImageProcessor, DetrForObjectDetection | |
import torch | |
from PIL import Image | |
import gradio as gr | |
# import requests | |
# import random | |
def detect_objects(image): | |
# Load the pre-trained DETR model | |
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-101") | |
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-101") | |
inputs = processor(images=image, return_tensors="pt") | |
outputs = model(**inputs) | |
# convert outputs (bounding boxes and class logits) to COCO API | |
# let's only keep detections with score > 0.9 | |
target_sizes = torch.tensor([image.size[::-1]]) | |
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0] | |
res = [] | |
for label in results["labels"]: | |
res.append(model.config.id2label[label.item()]) | |
return ','.join(res) | |
def upload_image(file): | |
image = Image.open(file.name) | |
image_with_boxes = detect_objects(image) | |
return image_with_boxes | |
iface = gr.Interface( | |
fn=upload_image, | |
inputs="file", | |
outputs="text", | |
title="Object Detection", | |
description="Upload an image and detect objects using DETR model.", | |
flagging_mode="never" | |
) | |
iface.launch() |