import os
import urllib.request
import numpy as np
import cv2
from ultralytics import YOLO
import gradio as gr

# --- Setup writable paths for YOLO ---
os.environ["YOLO_CONFIG_DIR"] = "/tmp"
os.environ["HOME"] = "/tmp"

# --- Download YOLOv8s weights if not already present ---
MODEL_PATH = "/tmp/yolov8s.pt"
if not os.path.exists(MODEL_PATH):
    print("Downloading YOLOv8s weights...")
    urllib.request.urlretrieve(
        "https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8s.pt",
        MODEL_PATH
    )

# --- Load the YOLOv8 model ---
model = YOLO(MODEL_PATH)

# --- Detection function ---
def detect_objects(image):
    """
    Input: image in BGR format (from Gradio/OpenCV)
    Output: annotated image, detected object names
    """
    # Convert BGR to RGB for YOLO
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Run YOLO inference directly on RGB image
    results = model(image_rgb, conf=0.25)  # confidence threshold 25%

    # Annotated image (YOLO expects RGB, but plot returns RGB)
    annotated_image = results[0].plot()

    # Convert annotated image back to BGR for OpenCV/Gradio display
    annotated_bgr = cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)

    # Extract detected object names
    if results[0].boxes is not None:
        detected_classes = [model.names[int(c)] for c in results[0].boxes.cls]
        detected_text = ", ".join(detected_classes) if detected_classes else "No objects detected"
    else:
        detected_text = "No objects detected"

    return annotated_bgr, detected_text

# --- Gradio Interface ---
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[
        gr.Image(type="numpy", label="Detected Objects"),
        gr.Textbox(label="Objects Detected")
    ],
    title="🧠 Object Detection App",
    description="Upload an image — YOLOv8s detects all objects and lists their names!"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)