Hugging Face's logo

Spaces:
Gradio-Blocks
/
Object-Detection-With-DETR-and-YOLOS
Runtime error

App Files Community
1
Object-Detection-With-DETR-and-YOLOS
File size: 3,538 Bytes 
69f28bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a801d9c
0ef0e61
e87525e
69f28bd
 
 
 
 
 
aa0f688
69f28bd
a801d9c
 
 
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
 
import io
import gradio as gr
import matplotlib.pyplot as plt
import requests, validators
import torch
from PIL import Image
from transformers import AutoFeatureExtractor, DetrForObjectDetection, YolosForObjectDetection

import os

# colors for visualization
COLORS = [
    [0.000, 0.447, 0.741],
    [0.850, 0.325, 0.098],
    [0.929, 0.694, 0.125],
    [0.494, 0.184, 0.556],
    [0.466, 0.674, 0.188],
    [0.301, 0.745, 0.933]
]

title = 'Object Detection App with DETR and YOLOS'

description = """
Links to HuggingFace Models:

- [facebook/detr-resnet-50](https://huggingface.co/facebook/detr-resnet-50)  
- [facebook/detr-resnet-101](https://huggingface.co/facebook/detr-resnet-101)  
- [hustvl/yolos-small](https://huggingface.co/hustvl/yolos-small)

"""

models = ["facebook/detr-resnet-50","facebook/detr-resnet-101",'hustvl/yolos-small']

options = gr.Dropdown(choices=models,label='Select Object Detection Model',show_label=True)

def make_prediction(img, feature_extractor, model):
    inputs = feature_extractor(img, return_tensors="pt")
    outputs = model(**inputs)
    img_size = torch.tensor([tuple(reversed(img.size))])
    processed_outputs = feature_extractor.post_process(outputs, img_size)
    return processed_outputs[0]

def fig2img(fig):
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    img = Image.open(buf)
    return img


def visualize_prediction(pil_img, output_dict, threshold=0.7, id2label=None):
    keep = output_dict["scores"] > threshold
    boxes = output_dict["boxes"][keep].tolist()
    scores = output_dict["scores"][keep].tolist()
    labels = output_dict["labels"][keep].tolist()
    if id2label is not None:
        labels = [id2label[x] for x in labels]

    plt.figure(figsize=(16, 10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=3))
        ax.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
    plt.axis("off")
    return fig2img(plt.gcf())

def detect_objects(model_name,url,image_upload,threshold):
    
    #Extract model and feature extractor
    feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
    
    if 'detr' in model_name:
        
        model = DetrForObjectDetection.from_pretrained(model_name)
        
    elif 'yolos' in model_name:
    
        model = YolosForObjectDetection.from_pretrained(model_name)
        
    if validators.url(url):
        image = Image.open(requests.get(url, stream=True).raw)
    elif image_upload:
        image = image_upload
    
    #Make prediction
    processed_outputs = make_prediction(image, feature_extractor, model)
    
    #Visualize prediction
    viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
    
    return viz_img   

examples= [['https://media-cldnry.s-nbcnews.com/image/upload/t_fit-1500w,f_auto,q_auto:best/newscms/2020_14/3290756/200331-wall-street-ew-343p.jpg',['IMG_5205.JPG'],[0.7]] 

gr.Interface(
    fn = detect_objects,
    inputs = [options,
             gr.Textbox(lines=1,label='Enter valid image URL here..'),
             gr.Image(type='pil'),
             gr.Slider(minimum=0.2,maximum=1,value=0.7,label='Prediction Threshold')],
    outputs=gr.Image(shape=(450,450)),
    title = title,
    description=description,
    examples=examples).launch()