import cv2 as cv
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download
from yolox import YoloX

# Download YOLOX model from Hugging Face (optional fallback)
model_path = hf_hub_download(
    repo_id="opencv/object_detection_yolox",
    filename="object_detection_yolox_2022nov.onnx"
)

# Initialize YOLOX model
model = YoloX(
    modelPath=model_path,
    confThreshold=0.5,
    nmsThreshold=0.5,
    objThreshold=0.5,
    backendId=cv.dnn.DNN_BACKEND_OPENCV,
    targetId=cv.dnn.DNN_TARGET_CPU
)

classes = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
           'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
           'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
           'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
           'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
           'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
           'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
           'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
           'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
           'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
           'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')

def letterbox(srcimg, target_size=(640, 640)):
    padded_img = np.ones((target_size[0], target_size[1], 3), dtype=np.float32) * 114.0
    ratio = min(target_size[0] / srcimg.shape[0], target_size[1] / srcimg.shape[1])
    resized_img = cv.resize(srcimg, (int(srcimg.shape[1] * ratio), int(srcimg.shape[0] * ratio)), interpolation=cv.INTER_LINEAR).astype(np.float32)
    padded_img[:int(srcimg.shape[0] * ratio), :int(srcimg.shape[1] * ratio)] = resized_img
    return padded_img, ratio

def unletterbox(bbox, scale):
    return bbox / scale

def visualize(dets, image, scale):
    res_img = image.copy()
    h, w = res_img.shape[:2]
    font_scale = max(0.5, min(w, h) / 640.0 * 0.5)
    thickness = max(1, int(font_scale * 2))

    for det in dets:
        box = unletterbox(det[:4], scale).astype(np.int32)
        score = det[-2]
        cls_id = int(det[-1])

        x0, y0, x1, y1 = box
        label = '{}:{:.1f}%'.format(classes[cls_id], score * 100)

        cv.rectangle(res_img, (x0, y0), (x1, y1), (0, 255, 0), thickness)
        (tw, th), _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
        cv.rectangle(res_img, (x0, y0), (x0 + tw + 2, y0 + th + 4), (255, 255, 255), -1)
        cv.putText(res_img, label, (x0, y0 + th), cv.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thickness)

    return res_img

def detect_objects(input_image):
    bgr = cv.cvtColor(input_image, cv.COLOR_RGB2BGR)
    input_blob, scale = letterbox(cv.cvtColor(bgr, cv.COLOR_BGR2RGB))

    results = model.infer(input_blob)
    if results is None or len(results) == 0:
        return input_image

    vis_image = visualize(results, bgr, scale)
    return cv.cvtColor(vis_image, cv.COLOR_BGR2RGB)

def clear_all():
    return None, None

def clear_output():
    return None

with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### YOLOX Object Detection (OpenCV + ONNX)")
    gr.Markdown("Upload an image to detect objects using YOLOX ONNX model and OpenCV DNN.")

    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload Image")
        output_image = gr.Image(type="numpy", label="Detected Objects")

    # Clear output when new image is uploaded
    image_input.change(fn=clear_output, outputs=[output_image])

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")

    submit_btn.click(fn=detect_objects, inputs=image_input, outputs=output_image)
    clear_btn.click(fn=clear_all, outputs=[image_input, output_image])

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/left.jpg"],
            ["examples/messi5.jpg"]
        ],
        inputs=image_input
    )


if __name__ == "__main__":
    demo.launch()