import torch
import json
from PIL import Image
from pathlib import Path
import gradio as gr
import numpy as np
import cv2
from ultralytics import YOLO

# 載入車框權重
car_mask_weights = 'mask.pt'
car_mask_model = torch.hub.load('./yolov5', 'custom', path=car_mask_weights, source="local")
car_mask_model.conf = 0.7  # 車框信心指數

# 載入車號權重
plate_weights = 'plate.pt'
plate_model = torch.hub.load('./yolov5', 'custom', path=plate_weights, source="local")
plate_model.conf = 0.5  # 車號信心指數
plate_model.classes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35]  # 過濾'-'符號

# 設定輸入和輸出資料夾
input_folder = 'input/'
output_folder = 'output/'

# 確保輸出資料夾存在
Path(output_folder).mkdir(parents=True, exist_ok=True)

# 將預測邏輯封裝成函數
def predict_license_plate(im):
    # 使用車框權重進行預測
    car_mask_results = car_mask_model(im, size=1280)

    # 取得車框預測的座標和類別
    car_mask_boxes = car_mask_results.pandas().xyxy[0]

    # 對每個車框進行車號預測
    plate_json = {}  # 將車號依序存成json
    for i, box in car_mask_boxes.iterrows():
        # 切割車框圖片
        x1, y1, x2, y2 = box['xmin'], box['ymin'], box['xmax'], box['ymax']

        # 放寬車牌裁切範圍
        expand_factor = 0.1  # 調整擴大比例
        expand_width = int((x2 - x1) * expand_factor)
        expand_height = int((y2 - y1) * expand_factor)
        x1 -= expand_width
        y1 -= expand_height
        x2 += expand_width
        y2 += expand_height

        # 限制裁切範圍在圖片邊界內
        x1 = max(x1, 0)
        y1 = max(y1, 0)
        x2 = min(x2, im.width)
        y2 = min(y2, im.height)

        car_mask_image = im.crop((x1, y1, x2, y2))

        # 使用車號權重進行預測
        plate_results = plate_model(car_mask_image, size=640)

        # 提取車號預測的結果
        plate_labels = plate_results.pandas().xyxy[0].sort_values('xmin')['name'].tolist()

        # 儲存圖片和預測結果
        plate_image_name = f'{Path(image_file).stem}_plate{i}.jpg'
        plate_image_path = Path(output_folder) / plate_image_name
        car_mask_image.save(plate_image_path)  # 儲存圖片

        licence_plate = ''.join(plate_labels)
        plate_json[plate_image_name] = licence_plate

    return [Image.fromarray(car_mask_results.ims[0]), plate_json]

# 將函數包裝成 Gradio 介面
inputs = gr.Image(type='pil', label="Original Image")
outputs = [gr.Image(type="pil", label="Output Image"),
           gr.JSON(label="Output JSON")]

title = "License_Plate_Prediction"
description = "Predict license plates using YOLOv5."
examples = [Image.open('input/1.jpg'), Image.open('input/2.jpg')]

gr.Interface(predict_license_plate, inputs, outputs, title=title, description=description, examples=examples).launch(enable_queue=True)