First iteration of the Gradio space

app.py

@@ -0,0 +1,73 @@
+import gradio as gr
+import spaces
+import torch
+
+from PIL import Image
+import requests
+from transformers import DetrImageProcessor
+from transformers import DetrForObjectDetection
+import matplotlib.pyplot as plt
+import io
+
+
+processor = DetrImageProcessor.from_pretrained("sergiopaniego/fashionpedia-finetuned_albumentations_coco")
+model = DetrForObjectDetection.from_pretrained("sergiopaniego/fashionpedia-finetuned_albumentations_coco")
+
+
+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]]
+
+def get_output_figure(pil_img, scores, labels, boxes):
+    plt.figure(figsize=(16, 10))
+    plt.imshow(pil_img)
+    ax = plt.gca()
+    colors = COLORS * 100
+    for score, label, (xmin, ymin, xmax, ymax), c in zip(scores.tolist(), labels.tolist(), boxes.tolist(), colors):
+        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=c, linewidth=3))
+        text = f'{model.config.id2label[label]}: {score:0.2f}'
+        ax.text(xmin, ymin, text, fontsize=15, 
+                bbox=dict(facecolor='yellow', alpha=0.5))
+    plt.axis('off')
+
+    return plt.gcf()
+
+@spaces.GPU
+def detect(image):
+    encoding = processor(image, return_tensors='pt')
+    print(encoding.keys())
+    
+    with torch.no_grad():
+        outputs = model(**encoding)
+
+    width, height = image.size
+    postprocessed_outputs = processor.post_process_object_detection(outputs, target_sizes=[(height, width)], threshold=0.9)
+    results = postprocessed_outputs[0]
+
+
+    output_figure = get_output_figure(image, results['scores'], results['labels'], results['boxes'])
+
+    buf = io.BytesIO()
+    output_figure.savefig(buf, bbox_inches='tight')
+    buf.seek(0)
+    output_pil_img = Image.open(buf)
+
+    return output_pil_img
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Object detection with DETR")
+    gr.Markdown(
+        """
+        This applciation uses DETR (DEtection TRansformers) to detect objects on images.
+        You can load an image and see the predictions for the objects detected along with the attention weights.
+        """
+    )
+    
+    gr.Interface(
+        fn=detect, 
+        inputs=gr.Image(label="Input image", type="pil"), 
+        outputs=[
+            gr.Image(label="Output prediction", type="pil")
+        ]
+    )#.launch()
+
+demo.launch(show_error=True)

requirements.txt

@@ -0,0 +1,3 @@
+transformers 
+timm
+torch