Use OpenCV

app.py

@@ -5,68 +5,38 @@ import time
 
 import torch
 import gradio as gr
-import matplotlib
-import matplotlib.pyplot as plt
-from PIL import Image
+import cv2
 
 from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 
 extractor = AutoFeatureExtractor.from_pretrained("hustvl/yolos-tiny")
 model = AutoModelForObjectDetection.from_pretrained("hustvl/yolos-tiny")
 
-matplotlib.pyplot.switch_backend('Agg') 
-
-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]
-]
-
+BBOX_COLOR = [255, 0, 0]
 PRED_THRESHOLD = 0.90
 
-def fig2img(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf)
-    buf.seek(0)
-    img = Image.open(buf)
-    return img
-
-def composite_predictions(img, processed_predictions):
-    keep = processed_predictions["labels"] == 1 # only interested in people 
-    boxes = processed_predictions["boxes"][keep].tolist()
-    scores = processed_predictions["scores"][keep].tolist()
-    labels = processed_predictions["labels"][keep].tolist()
-
-    labels = [model.config.id2label[x] for x in labels]
+def composite_predictions(img, processed_predictions, show_video=False):
+    interested_labels = processed_predictions["labels"] == 1 # only interested in people 
+    scores = processed_predictions["scores"][interested_labels].tolist()
+    boxes = [[int(j) for j in x] for x in processed_predictions["boxes"][interested_labels].tolist()]
+    labels = [model.config.id2label[x] for x in processed_predictions["labels"][interested_labels].tolist()]
     
-    plt.figure(figsize=(16, 10))
-    plt.imshow(img)
-    axis = plt.gca()
-    label_counts = defaultdict(int)
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, COLORS * len(boxes)):
-        axis.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=3))
-        axis.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
-        label_counts[label] += 1
-    plt.axis("off")
-    img = fig2img(plt.gcf())
-    matplotlib.pyplot.close()
-    return img, label_counts, datetime.datetime.now()
+    for score, box, label in zip(scores, boxes, labels):
+        cv2.rectangle(img, box, BBOX_COLOR, 1)
+        cv2.putText(img, f"{label}: {score:0.2f}", (box[0]+2, box[1]+10), cv2.FONT_HERSHEY_SIMPLEX, 0.33, BBOX_COLOR, 1, cv2.LINE_AA)
+    return img, len(boxes), datetime.datetime.now()
 
 def process(img):
-    #time.sleep(5)
     inputs = extractor(images=img, return_tensors="pt")
     outputs = model(**inputs)
-    img_size = torch.tensor([tuple(reversed(img.size))])
+    h, w, _ = img.shape
+    img_size = torch.tensor([(h, w)])
     processed = extractor.post_process_object_detection(outputs, PRED_THRESHOLD, img_size)
     
     # Composite image and prediction bounding boxes + labels prediction
     return composite_predictions(img, processed[0])
 
 
-
 with gr.Blocks() as demo:
     stream = gr.State()
     with gr.Row():
@@ -75,11 +45,9 @@ with gr.Blocks() as demo:
             attendance_label = gr.Label(label="Current Attendance")
         with gr.Row():
             with gr.Column(scale=1, min_width=600):
-                webcam = gr.Webcam(streaming=True, type='pil')
-                output = gr.Image(label="Composite")
-        webcam.stream(process, [webcam], [output, attendance_label, last_refresh_box])
-        #webcam.change(process, inputs=[], outputs=[output, gr.Label(label="Current Attendance"), last_refresh_box], max_batch_size=10, batch=True)
-    #demo.load(lambda: datetime.datetime.now(), None, last_refresh_box, every=10)
+                webcam = gr.Webcam(streaming=True)
+                output = gr.Image(label="Composite", visible=True)
+    webcam.stream(process, [webcam], [output, attendance_label, last_refresh_box])
 
 if __name__ == "__main__":
     demo.queue().launch()

requirements.txt

@@ -33,6 +33,7 @@ mdit-py-plugins==0.3.1
 mdurl==0.1.2
 multidict==6.0.2
 numpy==1.23.5
+opencv-python==4.6.0.66
 orjson==3.8.2
 packaging==21.3
 pandas==1.5.1