import os
import gradio as gr
from PIL import Image
import pytesseract
import yolov5

# load model
model = yolov5.load('keremberke/yolov5m-license-plate')

# set model parameters
model.conf = 0.5  # NMS confidence threshold
model.iou = 0.25  # NMS IoU threshold
model.agnostic = False  # NMS class-agnostic
model.multi_label = False  # NMS multiple labels per box
model.max_det = 1000  # maximum number of detections per image

def license_plate_detect(img):
    results = model(img, size=640)

    # parse results
    if len(results.pred):
        predictions = results.pred[0]
        boxes = predictions[:, :4] # x1, y1, x2, y2
        return boxes


def read_license_number(img):
    boxes = license_plate_detect(img)
    if len(boxes[0]):
        image = Image.fromarray(img)
        return [pytesseract.image_to_string(
                    image.crop(bbox.tolist()))
               for bbox in boxes]


def greet(img):
    boxes = license_plate_detect(img)
    image = Image.fromarray(img)
    r = 'greet'
    if len(boxes[0]):
        r = [pytesseract.image_to_string(
                    image.crop(bbox.tolist()))
               for bbox in boxes]
    return "Hello " + str(r) + "!!"


iface = gr.Interface(fn=greet, inputs="image", outputs="text")
iface.launch()