from PIL import Image
import gradio as gr
# import gradio as gr
import PIL.Image as Image
from ultralytics import ASSETS, YOLO

model_path = 'new_data_improved_object_detector.pt'
model = YOLO(model_path)

from PIL import Image
import gradio as gr

object_detector_model_path = 'new_data_improved_object_detector.pt'
# logo_detector_model_path ='logo_detector_grayscale_v2.pt'
logo_detector_model_path = 'logo_detector_june_1.pt'

object_model = YOLO(object_detector_model_path)
logo_model = YOLO(logo_detector_model_path)

def Get_logo_xywh(model_result_input):
    model_result = model_result_input[0]
    xywh = model_result.boxes.xywh.cpu().tolist()
    clss = model_result.boxes.cls.cpu().tolist()
    # names = model_result_input[0].names
    confidence = model_result.boxes.conf.cpu().tolist()
    xyxy = model_result.boxes.cpu().xyxy.tolist()

    return xywh, clss, confidence, xyxy

def predict_image(img, conf_threshold, iou_threshold):
    """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds."""

    resized_image = img.resize((640, 640))  
    # Convert the input image to grayscale
    img = resized_image.convert('L')

    logo_results = logo_model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )

    logo_im_arrays = []

    for r in logo_results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
        logo_im_arrays.append(im)

    object_result = object_model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )

    im_arrays = []

    for r in object_result:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
        im_arrays.append(im)


    logo_xywh, logo_clss,  logo_confidence, logo_xyxy = Get_logo_xywh(logo_results)
    object_xywh, object_clss,  object_confidence, object_xyxy = Get_logo_xywh(object_result)

    return logo_im_arrays,im_arrays, logo_xywh, logo_clss,  logo_confidence, logo_xyxy,object_xywh, object_clss,  object_confidence, object_xyxy

iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
    ],
    outputs=[
        gr.Gallery(label="logo Images"),
        gr.Gallery(label="object Images"),
        gr.JSON(label="Detection Bounding Boxes (l_xywh)"),
        gr.JSON(label="Detection Class Indices"),
        gr.JSON(label="Detection Confidence Scores"),
        gr.JSON(label="Detection Bounding Boxes (l_xyxy)"),
        gr.JSON(label="Detection Bounding Boxes (l_xywh)"),
        gr.JSON(label="Detection Class Indices"),
        gr.JSON(label="Detection Confidence Scores"),
        gr.JSON(label="Detection Bounding Boxes (l_xyxy)"),
    ],
    title="Ultralytics Gradio",
    description="Upload images for inference. The Ultralytics YOLOv8n model is used by default.",
)

if __name__ == "__main__":
    iface.launch(show_error=True)