from huggingface_hub import hf_hub_download
import gradio as gr
from image_inference import RTDETR_ONNX


model_path = hf_hub_download(
    repo_id="hasnatz/v-safe-rf-detr",
    filename="inference_model.onnx"
)

model = RTDETR_ONNX(model_path)

# Define inference function for Gradio
def predict(image, confidence=0.25, max_boxes=100):
    MAX_SIZE = 640   # you can tune this (e.g., 640, 1024)
    if max(image.size) > MAX_SIZE:
        image.thumbnail((MAX_SIZE, MAX_SIZE))
    # image is already a PIL.Image from Gradio
    return model.run_inference(image, confidence_threshold=confidence, max_number_boxes=max_boxes)

# Ready-made example images (local files or URLs)
examples = [
    ["examples/121113-F-LV838-027.jpg"],
    ["examples/goggles_bing_construction_goggles_000109.jpg"],
    ["examples/image-shows-busy-construction-site-where-concrete-mixer-truck-works-alongside-laborers-safety-gear-focus-teamwork-347908285 (Small).jpeg"],
    ["examples/istockphoto-1324894706-612x612.jpg"],
    ["examples/shutterstock_174689291.jpg"],
    ["examples/worker_bing_construction_building_worker_000043 (8).jpg"],
    ["examples/worker_bing_construction_building_worker_000066 (1).jpg"],
    ["examples/worker_bing_construction_building_worker_000091 (2).jpg"]
]


# Building Gradio UI


custom_theme = gr.themes.Base().set(
    body_background_fill="#0f0f11",   # background color
    block_background_fill="#0f0f11",  # blocks background
    block_border_color="#0f0f11",     # remove border feel
    background_fill_primary="#0f0f11" # for other sections
)


with gr.Blocks(theme=custom_theme) as demo:

    gr.HTML(
    """
    <div style="text-align: center;">
    <img 
        src='/gradio_api/file=Logo.png' 
        alt='My Image' 
        style='height: 100px; width: auto; display: block; margin: 0 auto;'
    >
    <h2>V-Safe: Construction Site Safety Detection Demo</h2>
    <br>
</div>
    """
    )

    

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Upload an Image", height=450)
            confidence = gr.Slider(0.0, 1.0, value=0.25, step=0.05, label="Confidence Threshold")
            
            run_btn = gr.Button("Run Inference")

        with gr.Column():
            output_img = gr.Image(type="pil", label="Annotated Result", height=450)

    run_btn.click(
        fn=predict,
        inputs=[input_img, confidence],
        outputs=output_img
    )

    gr.Examples(
        examples=examples,
        inputs=[input_img],
        outputs=output_img,
        fn=predict,
        cache_examples=True
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(allowed_paths=["Logo.png"])