import os
import torch

# CRITICAL: Redirect cache to temporary storage to avoid Hugging Face Space eviction
os.environ['TORCH_HOME'] = '/tmp/torch_cache'
os.environ['HUB_DIR'] = '/tmp/torch_hub'
os.environ['TMPDIR'] = '/tmp'
# Added to fix the Ultralytics config warning:
os.environ['YOLO_CONFIG_DIR'] = '/tmp/yolo_config' 

torch.hub.set_dir('/tmp/torch_hub')

import gradio as gr
from ultralytics import YOLO

# Load the model
model_path = "OceanCV_FirstPass.pt"
model = YOLO(model_path)

def run(image_path, conf, iou):
    # Predict using the slider values
    # Note: 'classes=0' is kept from your baseline template 
    results = model.predict(image_path, conf=conf, iou=iou, classes=0)
    
    # Reverse channels from BGR (OpenCV/YOLO default) to RGB (Gradio expectation)
    return results[0].plot()[:, :, ::-1]

title = "OceanCV First Pass Detector"
description = "Upload an image to detect objects using the first pass model."

# Build the interface with interactive sliders
interface = gr.Interface(
    fn=run,
    inputs=[
        gr.Image(type="filepath", label="Upload Image"),
        gr.Slider(minimum=0.05, maximum=1.0, value=0.20, step=0.05, label="Confidence Threshold"),
        gr.Slider(minimum=0.05, maximum=1.0, value=0.85, step=0.05, label="IoU Threshold")
    ],
    outputs=gr.Image(type="numpy", label="Detections"),
    title=title,
    description=description
)

# Launch the app with server_name and port explicitly set for HF Spaces
interface.queue().launch(server_name="0.0.0.0", server_port=7860)