# Are you wearing a mask?
import gradio as gr
import torch
import torchvision
import numpy as np
from PIL import Image

# Face masks
# TODO: Allow user selectable model?
model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/face_masks_v8.pt")




def yolo(im, size=640):
    g = (size / max(im.size))  # gain
    im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize

    results = model(im)  # inference
    results.render()  # updates results.imgs with boxes and labels
    return Image.fromarray(results.imgs[0])


inputs = gr.inputs.Image(type='pil', label="Original Image")
outputs = gr.outputs.Image(type="pil", label="Output Image")

title = "Are you wearing a mask?"
description = "Detecting masked and unmasked faces with YOLOv5. Take a picture, upload an image, or click an example image to use."
article = "<p style='text-align: center'>This app makes predictions using a YOLOv5s model that was fine tuned on a <a href='https://www.kaggle.com/datasets/henrylydecker/face-masks'>dataset</a> of people with and without masks. All of the code  for training the model is available on <a href='https://github.com/hlydecker/are-you-wearing-a-mask'>GitHub</a>. This app and the model behind it were created by Henry Lydecker, for a <a href='https://github.com/Sydney-Informatics-Hub/cv-demo'>course</a> he developed for the Sydney Informatics Hub, a Core Research Facility of The University of Sydney. Find out more about the YOLO model from the original creator, <a href='https://pjreddie.com/darknet/yolo/'>Joseph Redmon</a>. YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset and developed by Ultralytics, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"

examples = [['data/picard.jpg'], ['data/crowd.jpeg'],['data/baseball2.jpeg'],['data/santa-claus-orig.jpg'],['data/kfc_anime2.jpg'],['data/doge2.webp'],['data/cat_mask.jpg']]
gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(enable_queue=True)