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-50") | |
| model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") | |
| 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() |