from transformers import pipeline

from PIL import Image, ImageDraw

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