hustvl/yolos-base obj detection

app.py

@@ -0,0 +1,11 @@
+import gradio as gr
+
+from ppl_detector import detect
+
+with gr.Blocks() as demo:
+    img_input = gr.Image(type="filepath", label="Input Image")
+    img_output = gr.Image(label="Output Image with Detections")
+    
+    img_input.change(fn=detect, inputs=img_input, outputs=img_output)
+       
+demo.launch()

ppl_detector.py

@@ -0,0 +1,40 @@
+from transformers import pipeline
+
+from PIL import Image, ImageDraw
+from pprint import pprint
+
+detector = pipeline("object-detection", model="hustvl/yolos-base")
+source = None
+
+def draw_bbox(bboxes):
+    global source
+    if source is None:
+        return
+    draw = ImageDraw.Draw(source)
+    for bb in bboxes:
+        draw.rectangle(bb, outline='yellow', width=2)
+
+
+def detect(img):
+    global source
+    if source is None and img is not None:
+        source = Image.open(img)
+    else:
+        return None
+    
+    objects = detector(source)
+    persons = [
+        o for o in objects 
+        if o['label'] == 'person' and o['score'] > 0.83]
+    bboxes = list(map(lambda p: (p['box']['xmin'], p['box']['ymin'], p['box']['xmax'], p['box']['ymax']), persons))
+    n = len(persons)
+    
+    print(f"it's got {n} {'ppl' if n > 1 else 'person'} in the image")
+
+    draw_bbox(bboxes)
+    return source
+
+
+
+
+