WIP

app.py

@@ -32,7 +32,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()
@@ -49,7 +49,7 @@ def composite_predictions(img, processed_predictions):
     plt.axis("off")
     img = fig2img(plt.gcf())
     matplotlib.pyplot.close()
-    return img
+    return img, len(boxes)
 
 def process(img):
     inputs = extractor(images=img, return_tensors="pt")
@@ -60,5 +60,7 @@ 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, type='pil')], outputs=["image"], live=True)
-demo.launch()
+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)
+
+if __name__ == "__main__":
+    demo.launch()