App bug fix.

app.py

@@ -11,7 +11,8 @@ import pytorch_lightning as pl
 
 from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 
-from PIL import Image
+from PIL import Image, ImageDraw
+import cv2
 import matplotlib.pyplot as plt
 
 id2label = {1: 'person', 2: 'rider', 3: 'car', 4: 'bus', 5: 'truck', 6: 'bike', 7: 'motor', 8: 'traffic light', 9: 'traffic sign', 10: 'train'}
@@ -100,21 +101,23 @@ def rescale_bboxes(out_bbox, size):
     return b
 
 def plot_results(pil_img, prob, boxes):
-    fig = plt.figure(figsize=(16,10), dpi=120)
-    plt.imshow(pil_img)
-    ax = plt.gca()
+
+    draw = ImageDraw.Draw(pil_img)
     colors = COLORS * 100
+    
     for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
         cl = p.argmax()
         c = colors[cl]
-        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
-                                   fill=False, color=c, linewidth=2))
-        text = f'{id2label[cl.item()]}: {p[cl]:0.2f}'
-        ax.text(xmin, ymin, text, fontsize=10,
-                bbox=dict(facecolor=c, alpha=0.5))
-    plt.axis('off')
-    # return Image.frombytes('RGB', fig.canvas.get_width_height(),fig.canvas.tostring_rgb())
-    return fig
+
+        draw.rectangle([xmin, ymin, xmax - xmin, ymax - ymin], outline=c, width=2)
+        draw.text(
+            [xmin + 5, ymin + 5], 
+            f'{id2label[cl.item()]}: {p[cl]:0.2f}',
+            fill=c)
+        # ax.text(xmin, ymin, text, fontsize=10,
+        #         bbox=dict(facecolor=c, alpha=0.5))
+    return Image.fromarray(pil_img[:,:,::-1])
+    # return fig
 
 
 def generate_preds(processor, model, image):
@@ -145,8 +148,8 @@ def detect(img):
 interface = gr.Interface(
     fn=detect,
     inputs=[gr.Image(type="pil")], 
-    # outputs=gr.Image(type="pil"),
-    outputs = ['plot'],
+    outputs=gr.Image(type="pil"),
+    # outputs = ['plot'],
     examples=[["./imgs/example1.jpg"], ["./imgs/example2.jpg"]],
     title="YOLOS for traffic object detection",
     description="A downstream application for <a href='https://huggingface.co/docs/transformers/model_doc/yolos' style='text-decoration: underline' target='_blank'>YOLOS</a> which can performe traffic object detection. ")

requirements.txt

@@ -5,4 +5,5 @@ git+https://github.com/huggingface/transformers.git@main
 matplotlib>=3.2.2
 Pillow>=7.1.2
 torch==1.10.0
-pytorch-lightning==1.9.3
+pytorch-lightning==1.9.3
+opencv-python>=4.1.1