app.py

"""
DEBRI-X Gradio Space Demo
Drop this file + yolo_best.pt into a Hugging Face Space (SDK=gradio).
requirements.txt: ultralytics opencv-python-headless huggingface_hub gradio
"""
import gradio as gr
import numpy as np, cv2
from PIL import Image

MODEL_REPO = "Premchan369/debrex-space-debris-capture"
try:
    from huggingface_hub import hf_hub_download
    from ultralytics import YOLO
    wt   = hf_hub_download(MODEL_REPO, "yolo_best.pt")
    yolo = YOLO(wt)
    LOADED = True
except Exception as e:
    print(f"Model load failed: {e}")
    LOADED = False

def detect(image, conf=0.25):
    if image is None:
        return None, "Please upload an image."
    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    if not LOADED:
        H, W = img_bgr.shape[:2]
        cv2.rectangle(img_bgr, (W//4, H//4), (3*W//4, 3*H//4), (0,200,100), 2)
        return Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)), "Demo mode."
    yolo.conf = conf
    res = yolo(img_bgr, verbose=False)
    ann = res[0].plot()
    n   = len(res[0].boxes)
    cfs = [f"{c:.2f}" for c in res[0].boxes.conf.cpu().numpy()]
    info = f"Objects: {n}\nConf: {', '.join(cfs) or 'none'}\nModel: YOLOv8n (DEBRI-X)"
    return Image.fromarray(cv2.cvtColor(ann, cv2.COLOR_BGR2RGB)), info

with gr.Blocks(title="DEBRI-X Debris Detection") as demo:
    gr.Markdown("# 🛰️ DEBRI-X — Space Debris Detection")
    with gr.Row():
        with gr.Column():
            inp  = gr.Image(type="pil", label="Input")
            conf = gr.Slider(0.05, 0.95, 0.25, label="Confidence")
            btn  = gr.Button("Detect", variant="primary")
        with gr.Column():
            out_img  = gr.Image(type="pil", label="Result")
            out_info = gr.Textbox(label="Info", lines=4)
    btn.click(detect, [inp, conf], [out_img, out_info])

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