Updated app

app.py

@@ -18,22 +18,17 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 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_figure(in_pil_img, in_results):
+def get_output_figure(pil_img, scores, labels, boxes):
     plt.figure(figsize=(16, 10))
-    plt.imshow(in_pil_img)
-    #pyplot.gcf()
+    plt.imshow(pil_img)
     ax = plt.gca()
-
-    for prediction in in_results:
-        selected_color = choice(COLORS)
-
-        x, y = prediction['box']['xmin'], prediction['box']['ymin'],
-        w, h = prediction['box']['xmax'] - prediction['box']['xmin'], prediction['box']['ymax'] - prediction['box']['ymin']
-
-        ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3))
-        ax.text(x, y, f"{prediction['label']}: {round(prediction['score']*100, 1)}%", fontdict=fdic)
-
-    plt.axis("off")
+    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()
 
@@ -46,9 +41,14 @@ def detect(image):
     with torch.no_grad():
         outputs = model(**encoding)
 
-    print(outputs)
+    #print(outputs)
+    width, height = image.size
+    postprocessed_outputs = processor.post_process_object_detection(outputs, target_sizes=[(height, width)], threshold=0.9)
+    results = postprocessed_outputs[0]
+
+    print(results)
 
-    output_figure = get_figure(image, outputs)
+    output_figure = get_output_figure(image, results['scores'], results['labels'], results['boxes'])
 
     buf = io.BytesIO()
     output_figure.savefig(buf, bbox_inches='tight')