import gradio as gr
import cv2
import joblib
from skimage.feature import hog
import numpy as np

MODEL_PATH = './hog_lreg_model_4.pkl'
clf = joblib.load(MODEL_PATH)

def ocr(pil_image):
    # Convert PIL image to OpenCV format
    im = np.array(pil_image)
    im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)

    im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    ret, im_th = cv2.threshold(im_gray, 120, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    ctrs, hier = cv2.findContours(im_th, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    bboxes = [cv2.boundingRect(c) for c in ctrs]
    sorted_bboxes = sorted(bboxes, key=lambda b: b[0])  # Sort by x-coordinate

    plate_char = []
    image_height, image_width = im.shape[:2]
    height_threshold = image_height * 0.3
    width_threshold = image_width * 0.3

    for num, i_bboxes in enumerate(sorted_bboxes):
        [x, y, w, h] = i_bboxes
        if h > height_threshold and w < width_threshold:
            roi = im_gray[y:y + h, x:x + w]
            roi = cv2.resize(roi, (64, 128), interpolation=cv2.INTER_AREA)
            roi_hog_fd = hog(roi, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(1, 1))
            nbr = clf.predict(np.array([roi_hog_fd]))
            plate_char.append(str(nbr[0]))

    return ''.join(plate_char)

# Create Gradio interface
interface = gr.Interface(
    fn=ocr,
    inputs=gr.Image(type="pil", label="Upload License Plate Image"),
    outputs=gr.Textbox(label="Predicted License Plate"),
    title="Automatic License Plate Recognition",
    description="Upload an image of a license plate, and the system will predict the text on it.",
)

# Launch the Gradio app
interface.launch()