jardin/test_yolo/test_yolo_leaf.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path


TOKEN_FR = {
    "healthy": "saine",
    "leaf": "feuille",
    "plant": "plante",
    "disease": "maladie",
    "blight": "mildiou",
    "rust": "rouille",
    "spot": "tache",
    "scab": "tavelure",
    "mold": "moisissure",
    "mildew": "oidium",
    "powdery": "poudreux",
    "bacterial": "bacterien",
    "viral": "viral",
    "virus": "virus",
    "yellow": "jaune",
    "mosaic": "mosaique",
    "late": "tardif",
    "early": "precoce",
    "septoria": "septoriose",
    "target": "cible",
    "black": "noire",
    "rot": "pourriture",
}


def translate_label_to_fr(label: str) -> str:
    tokens = (
        label.replace("-", " ")
        .replace("_", " ")
        .replace(",", " ")
        .replace("(", " ")
        .replace(")", " ")
        .split()
    )
    if not tokens:
        return label
    translated = [TOKEN_FR.get(tok.lower(), tok.lower()) for tok in tokens]
    return " ".join(translated)


def resolve_model(model_arg: str) -> tuple[str, str]:
    model_path = Path(model_arg)
    if model_path.exists():
        return str(model_path), model_arg

    lower = model_arg.lower()
    if "/" in model_arg and not lower.endswith((".pt", ".onnx", ".engine", ".torchscript")):
        try:
            from huggingface_hub import hf_hub_download
        except Exception as exc:
            raise SystemExit(
                "Le modèle semble être un repo Hugging Face. Installe:\n"
                "pip install huggingface_hub\n"
                f"Détail: {exc}"
            )

        try:
            downloaded = hf_hub_download(repo_id=model_arg, filename="best.pt")
            return downloaded, model_arg
        except Exception as exc:
            raise SystemExit(
                f"Impossible de télécharger best.pt depuis le repo Hugging Face '{model_arg}'.\n"
                f"Détail: {exc}"
            )

    return model_arg, model_arg


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Test YOLO feuille/plante avec sortie JSON et image annotée."
    )
    parser.add_argument("--image", required=True, help="Chemin image à analyser")
    parser.add_argument(
        "--model",
        default="foduucom/plant-leaf-detection-and-classification",
        help="Chemin vers .pt/.onnx ou repo Hugging Face (best.pt)",
    )
    parser.add_argument("--conf", type=float, default=0.25, help="Seuil confiance")
    parser.add_argument("--iou", type=float, default=0.45, help="Seuil IoU NMS")
    parser.add_argument("--max-det", type=int, default=1000, help="Nombre max de détections")
    parser.add_argument(
        "--json-out",
        default="test_yolo/output/detections.json",
        help="Fichier JSON de sortie",
    )
    parser.add_argument(
        "--image-out",
        default="test_yolo/output/annotated.jpg",
        help="Image annotée de sortie",
    )
    args = parser.parse_args()

    image_path = Path(args.image)
    if not image_path.exists():
        raise SystemExit(f"Image introuvable: {image_path}")

    try:
        import cv2
        from ultralytics import YOLO
    except Exception as exc:
        raise SystemExit(
            "Dépendances manquantes. Installe dans un venv:\n"
            "pip install ultralytics opencv-python matplotlib huggingface_hub\n"
            f"Détail: {exc}"
        )

    model_source, model_display = resolve_model(args.model)
    model = YOLO(model_source)
    model.overrides["conf"] = args.conf
    model.overrides["iou"] = args.iou
    model.overrides["max_det"] = args.max_det

    results = model.predict(str(image_path))
    result = results[0]

    boxes = result.boxes
    names = getattr(result, "names", {})
    detections: list[dict] = []
    for i, box in enumerate(boxes):
        cls_id = int(box.cls.item())
        conf = float(box.conf.item())
        xyxy = [float(v) for v in box.xyxy[0].tolist()]
        detections.append(
            {
                "index": i,
                "class_id": cls_id,
                "class_name": names.get(cls_id, str(cls_id)),
                "class_name_fr": translate_label_to_fr(names.get(cls_id, str(cls_id))),
                "confidence": conf,
                "bbox_xyxy": xyxy,
            }
        )

    payload = {
        "image": str(image_path),
        "model": model_display,
        "model_source": model_source,
        "params": {"conf": args.conf, "iou": args.iou, "max_det": args.max_det},
        "count": len(detections),
        "detections": detections,
    }

    json_out = Path(args.json_out)
    json_out.parent.mkdir(parents=True, exist_ok=True)
    json_out.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")

    annotated = result.plot()
    image_out = Path(args.image_out)
    image_out.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(image_out), annotated)

    print(f"Détections: {len(detections)}")
    print(f"JSON: {json_out}")
    print(f"Image annotée: {image_out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())