WIP

app.py

@@ -1,4 +1,7 @@
+from collections import defaultdict
+import datetime
 import io
+import time
 
 import torch
 import gradio as gr
@@ -15,11 +18,11 @@ 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]
+    # [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]
 ]
 
 PRED_THRESHOLD = 0.90
@@ -32,7 +35,7 @@ def fig2img(fig):
     return img
 
 def composite_predictions(img, processed_predictions):
-    keep = processed_predictions["labels"] >= 1 # only interested in people 
+    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()
@@ -42,16 +45,18 @@ def composite_predictions(img, processed_predictions):
     plt.figure(figsize=(16, 10))
     plt.imshow(img)
     axis = plt.gca()
-    colors = COLORS * 100
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
+    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, len(boxes)
+    return img, label_counts, 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))])
@@ -60,7 +65,51 @@ def process(img):
     # Composite image and prediction bounding boxes + labels prediction
     return composite_predictions(img, processed[0])
 
-demo = gr.Interface(fn=process, inputs=[gr.Image(source="webcam", streaming=True, tool=None, type='pil')], outputs=["image", gr.Label(label="Current Attendance")], live=True)
+
+
+with gr.Blocks() as demo:
+    stream = gr.State()
+    with gr.Row():
+        with gr.Column(scale=1, min_width=600):
+            last_refresh_box = gr.Textbox(label="Last updated")
+            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)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()
+
+
+
+# import gradio as gr
+# import numpy as np
+# import time
+
+# def add_to_stream(audio, instream):
+#     time.sleep(1)
+#     if audio is None:
+#         return gr.update(), instream
+#     if instream is None:
+#         ret = audio
+#     else:
+#         ret = (audio[0], np.concatenate((instream[1], audio[1])))
+#     return ret, ret
+
+
+# with gr.Blocks() as demo:
+#     inp = gr.Audio(source="microphone")
+#     out = gr.Audio()
+#     stream = gr.State()
+#     clear = gr.Button("Clear")
+
+#     inp.stream(add_to_stream, [inp, stream], [out, stream])
+#     clear.click(lambda: [None, None, None], None, [inp, out, stream])
+
+
+# if __name__ == "__main__":
+#     demo.launch()