import gradio as gr
from fastai.vision.all import *


def get_headcount(filepath):
  filepath = str(filepath)
  filename = filepath.split('/')[-1]

  return df[df['Name'] == filename]['HeadCount'].values[0]


learn = load_learner('export_ver2.pkl')

def predict(img):
  img = PILImage.create(img)
  op = learn.predict(img)
  return int(op[0][0])




title = "Face count"
description = "This model that is trained to counts the number of faces  in the uploaded picture."

eg = ['cam_mit.jpg', 'dunphys.jpg', 'group.jpg']

gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(512, 512)),
    outputs=gr.outputs.Textbox(type="number", label="Number of faces"),
    title=title,
    description=description,
    examples=eg,
).launch(share=False)