import gradio as gr
from crowd_counter import CrowdCounter
from PIL import Image
import cv2

crowd_counter_engine = CrowdCounter()

def predict(inp):
    inp = Image.fromarray(inp.astype('uint8'), 'RGB')
    response = crowd_counter_engine.inference(inp)
    crowd_count = response[0]
    pred_img= response[1]
    return cv2.cvtColor(pred_img, cv2.COLOR_BGR2RGB), crowd_count

title = "Crowd Counter Demo"
desc = "A Demo of Proposal Point Prediction for Crowd Counting - Powered by P2PNet"
examples = [
    ["images/img-1.jpg"],
    ["images/img-2.jpg"],
    ["images/img-3.jpg"],
]
inputs = gr.inputs.Image(label="Image of Crowd")
outputs = [gr.outputs.Image(label="Proposal Points Prediction",type = "numpy"), gr.outputs.Label(label="Predicted Count",type = "numpy")]
gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=title, description=desc, examples=examples,
             allow_flagging=False).launch()