ppl_detector.py

from transformers import pipeline
from PIL import Image, ImageDraw
from pprint import pprint

detector = pipeline("object-detection", model="hustvl/yolos-base")
source = None
state = 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, result
    if img is not None:
        source = Image.open(img)
    else:
        return None
    
    print("processing img...")
    
    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)
    result = f"it's got {n} {'ppl' if n > 1 else 'person'} in the image"

    draw_bbox(bboxes)
    return source

def set_result(*args):
    global result
    if result is None:
        return f"<h1 id='result' style='display: flex;justify-content: center;'>it counts ppl, upload an image, see how it goes</h1>"
    else:
        title = f"<h1 id='result' style='display: flex;justify-content: center;'>{result}</h1>"
        result = None
        return title