import gradio as gr
import torch
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
import cv2
import numpy as np

model = torch.hub.load('ultralytics/yolov5', 'custom', path='model/yolov5n_rebar_kaggle.pt')

def yolo(im, conf, iou, size=640):

    mask = np.array(im["mask"])
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
      mask = np.zeros(mask.shape, np.uint8)
      cnt = contours[0]
      mask = cv2.drawContours(mask, [cnt], 0, 255, -1)
      im = np.array(im["image"])
      im = cv2.bitwise_and(im, im, mask=mask)
      im = Image.fromarray(im)
    else:
      im = im["image"]
    
    model.conf = conf
    model.iou = iou
    results = model(im, size=size)  # custom inference size
     # inference

    # output_im = Image.fromarray(results.render(labels=False)[0])
    # output_im = results.render(labels=False)[0]
    output_im = np.array(im)

    pred = results.pandas().xyxy[0]
    counting = pred.shape[0]
    text = f"{counting} objects"
    for index, row in pred.iterrows():
      cv2.circle(output_im, (int((row["xmin"] + row["xmax"]) * 0.5), int((row["ymin"] + row["ymax"]) * 0.5)), int((row["xmax"] - row["xmin"]) * 0.5 * 0.6), (255, 0, 0), -1)
 
    return Image.fromarray(output_im), text

slider_step = 0.05
nms_conf = 0.25
nms_iou = 0.1
# inputs_image = gr.inputs.Image(type='pil', label="Original Image")
inputs_image = gr.inputs.Image(tool="sketch", label="Original Image",type="pil")
inputs_conf = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="Conf Thres")
inputs_iou = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU Thres")

inputs = [inputs_image, inputs_conf, inputs_iou]

outputs_image = gr.outputs.Image(type="pil", label="Output Image")
outputs_text = gr.Textbox(label="Number of objects")
outputs = [outputs_image, outputs_text]

title = "OBJECT COUNTING"
description = "Object counting demo. Upload an image or click an example image to use. You can select the area to count by drawing a closed area on the input image."
article = "<p style='text-align: center'>Counting objects in image</a></p>"
examples = [['./images/S__275668998.jpg'], ['./images/S__271433737.jpg']]

gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, cache_examples=False, analytics_enabled=False).launch(
    debug=True)#, share=True)