app.py

import os
import urllib.request
import cv2
import numpy as np
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 (BGR)
    Output: annotated image, detected object names
    """
    # Convert BGR (OpenCV) to RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Resize for consistent detection
    image_resized = cv2.resize(image_rgb, (640, 640))

    # Run YOLO inference with confidence threshold 0.2
    results = model(image_resized, conf=0.2)

    # Annotated image
    annotated_image = results[0].plot()

    # Extract detected class 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_image, 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!"
)

# --- Launch the app ---
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)