app.py

import os, tempfile, cv2, numpy as np, gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

REPO_ID = "harpreetsahota/car-dd-segmentation-yolov11"
WEIGHT_FILE = "best.pt"
weights_path = hf_hub_download(repo_id=REPO_ID, filename=WEIGHT_FILE)
model = YOLO(weights_path)  # charge YOLO11-seg (segmentation d'instances)

def predict_image(image, conf=0.5, imgsz=1024):
    if image is None:
        raise ValueError("Aucune image fournie")

    # si l'image arrive en numpy, on force le bon type
    import numpy as np
    if isinstance(image, np.ndarray):
        image = image.astype("uint8")

    results = model.predict(
        source=image,
        conf=float(conf),
        imgsz=int(imgsz)
    )

    if len(results) == 0:
        raise RuntimeError("Aucun résultat retourné par le modèle")

    return results[0].plot()

def predict_video(video, conf=0.5, imgsz=1024, fps_out=15):
    cap = cv2.VideoCapture(video)
    if not cap.isOpened(): raise RuntimeError("Impossible d'ouvrir la vidéo")
    frames = []
    input_fps = cap.get(cv2.CAP_PROP_FPS) or 25
    step = max(1, int(round(input_fps / fps_out)))
    idx = 0
    while True:
        ret, frame = cap.read()
        if not ret: break
        if idx % step != 0: idx += 1; continue
        idx += 1
        res = model(frame, conf=conf, imgsz=int(imgsz))
        frames.append(res[0].plot())
    cap.release()
    if not frames: return None
    h, w = frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    vw = cv2.VideoWriter(tmp.name, fourcc, fps_out, (w, h))
    for f in frames: vw.write(f)
    vw.release()
    return tmp.name

with gr.Blocks() as demo:
    gr.Markdown("### Car damage segmentation (Ultralytics YOLO-seg, CarDD)")
    with gr.Tab("Image"):
        i = gr.Image(type="numpy", label="Photo véhicule")
        c = gr.Slider(0.1, 0.9, 0.5, step=0.05, label="Seuil (conf)")
        s = gr.Slider(640, 1280, 1024, step=64, label="imgsz")
        o = gr.Image(type="numpy", label="Résultat")
        gr.Button("Lancer").click(predict_image, [i, c, s], o)
    with gr.Tab("Vidéo"):
        vin = gr.Video(label="MP4")
        c2 = gr.Slider(0.1, 0.9, 0.5, step=0.05, label="Seuil (conf)")
        s2 = gr.Slider(640, 1280, 1024, step=64, label="imgsz")
        vout = gr.Video(label="Vidéo annotée")
        gr.Button("Lancer").click(predict_video, [vin, c2, s2], vout)
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)