[ADD] adding a IoU.py, a script to calculate the IoU metric between 2 yolo models.

.gitignore

@@ -31,4 +31,5 @@ lib64/
 generated_html/
 annotations/
 dataset/
-test/
+test/
+yolov5/

IoU_yolo.py

@@ -0,0 +1,110 @@
+import os
+import numpy as np
+
+GT_DIR = "test_images/wit396_pdf545/labels"
+PRED_DIR = ""
+ 
+def yolo_to_corners(box):
+    """Convertit (x_center, y_center, w, h) en (x1, y1, x2, y2)."""
+    x_c, y_c, w, h = box
+    x1 = x_c - w / 2
+    y1 = y_c - h / 2
+    x2 = x_c + w / 2
+    y2 = y_c + h / 2
+    return np.array([x1, y1, x2, y2])
+ 
+ 
+def get_iou(box_a, box_b):
+    """Calcule l'IoU entre deux boxes au format YOLO (x_center, y_center, w, h)."""
+    a = yolo_to_corners(box_a)
+    b = yolo_to_corners(box_b)
+ 
+    ix1 = np.maximum(a[0], b[0])
+    iy1 = np.maximum(a[1], b[1])
+    ix2 = np.minimum(a[2], b[2])
+    iy2 = np.minimum(a[3], b[3])
+ 
+    i_width = np.maximum(ix2 - ix1, 0.0)
+    i_height = np.maximum(iy2 - iy1, 0.0)
+    area_of_intersection = i_width * i_height
+ 
+    area_a = max(a[2] - a[0], 0.0) * max(a[3] - a[1], 0.0)
+    area_b = max(b[2] - b[0], 0.0) * max(b[3] - b[1], 0.0)
+ 
+    area_of_union = area_a + area_b - area_of_intersection
+    if area_of_union <= 0:
+        return 0.0
+ 
+    return float(area_of_intersection / area_of_union)
+ 
+ 
+def parse_yolo_file(filepath):
+    """Lit un .txt YOLO et retourne une liste de {class_id, box}."""
+    boxes = []
+    if not os.path.exists(filepath):
+        return boxes
+ 
+    with open(filepath, "r") as f:
+        for line in f:
+            parts = line.strip().split()
+            if not parts:
+                continue
+            class_id = int(float(parts[0]))
+            xc, yc, w, h = (float(p) for p in parts[1:5])
+            boxes.append({"class_id": class_id, "box": [xc, yc, w, h]})
+    return boxes
+ 
+ 
+def match_boxes(gt_boxes, pred_boxes):
+    """Associe GT et prédictions par classe, en priorisant le meilleur IoU."""
+    matches = []
+    unmatched_gt = list(gt_boxes)
+    unmatched_pred = list(pred_boxes)
+ 
+    classes = set(b["class_id"] for b in gt_boxes) | set(b["class_id"] for b in pred_boxes)
+ 
+    for cls in classes:
+        gt_cls = [b for b in unmatched_gt if b["class_id"] == cls]
+        pred_cls = [b for b in unmatched_pred if b["class_id"] == cls]
+ 
+        pairs = []
+        for gt_b in gt_cls:
+            for pred_b in pred_cls:
+                iou = get_iou(gt_b["box"], pred_b["box"])
+                pairs.append((iou, gt_b, pred_b))
+ 
+        pairs.sort(key=lambda x: x[0], reverse=True)
+        used_gt, used_pred = set(), set()
+        for iou, gt_b, pred_b in pairs:
+            if id(gt_b) in used_gt or id(pred_b) in used_pred:
+                continue
+            matches.append((gt_b, pred_b, iou))
+            used_gt.add(id(gt_b))
+            used_pred.add(id(pred_b))
+            unmatched_gt.remove(gt_b)
+            unmatched_pred.remove(pred_b)
+ 
+    return matches, unmatched_gt, unmatched_pred
+ 
+gt_files = sorted(f for f in os.listdir(GT_DIR) if f.endswith(".txt"))
+ 
+all_ious = []
+ 
+for filename in gt_files:
+    gt_path = os.path.join(GT_DIR, filename)
+    pred_path = os.path.join(PRED_DIR, filename)
+ 
+    gt_boxes = parse_yolo_file(gt_path)
+    pred_boxes = parse_yolo_file(pred_path)
+ 
+    matches, unmatched_gt, unmatched_pred = match_boxes(gt_boxes, pred_boxes)
+    ious = [m[2] for m in matches]
+    all_ious.extend(ious)
+ 
+    mean_iou = np.mean(ious) if ious else 0.0
+    print(f"{filename}: {len(matches)} paires, IoU moyen = {mean_iou:.4f}, "
+          f"GT manqués = {len(unmatched_gt)}, faux positifs = {len(unmatched_pred)}")
+ 
+print(f"Fichiers traités : {len(gt_files)}")
+print(f"Paires appariées : {len(all_ious)}")
+print(f"IoU moyen global : {np.mean(all_ious) if all_ious else 0.0:.4f}")

inference.py

@@ -1,14 +1,14 @@
 import argparse
 import glob
 import os
-import matplotlib.pyplot as plt
- 
 from ultralytics import YOLO
 
-# uv run predict.py  path_to_model.pt  path_to_images ./test
+# uv run predict.py path_to_model.pt path_to_images ./test
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Prédiction YOLO sur un dossier d'images avec sauvegarde des résultats annotés."
+        description="Prédiction YOLO sur un dossier d'images avec sauvegarde des résultats annotés et des labels .txt"
     )
 
     parser.add_argument(
@@ -24,41 +24,62 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "output_path",
         type=str,
-        help="Dossier de sortie pour les images annotées",
+        help="Dossier de sortie pour les images annotées et les labels",
     )
 
     return parser.parse_args()
 
 
-def predict(model_path:str, images_path:str, output_path:str):
+def save_yolo_txt(result, txt_path):
+    """
+    Sauvegarde les prédictions d'un résultat Ultralytics au format YOLO :
+    <class_id> <x_center> <y_center> <width> <height> <confidence>
+    (coordonnées normalisées entre 0 et 1)
+    """
+    boxes = result.boxes
+    with open(txt_path, "w") as f:
+        for box in boxes:
+            class_id = int(box.cls[0])
+            conf = float(box.conf[0])
+            x_center, y_center, width, height = box.xywhn[0].tolist()
+            f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {conf:.6f}\n")
+
+
+def predict(model_path: str, images_path: str, output_path: str):
     model = YOLO(model_path)
     images = glob.glob(os.path.join(images_path, "*.jpg"))
     selection = images[:25]
-    print(selection)
+    print(f"{len(selection)} images sélectionnées pour la prédiction")
+
     results = model(selection)
-    name_model = model_path.split('/')
 
-    output_dir = output_path
-    model_dir = name_model[0]
+    # Nom du modèle (sans extension) pour organiser les résultats par modèle testé
+    model_name = model_path.split('/')
+    model_name = model_name[0]
+
+    images_dir = os.path.join(output_path, model_name, "images")
+    labels_dir = os.path.join(output_path, model_name, "labels")
+    os.makedirs(images_dir, exist_ok=True)
+    os.makedirs(labels_dir, exist_ok=True)
 
-    if not os.path.isdir(output_dir):
-        os.mkdir(output_dir)
+    for image_path, result in zip(selection, results):
+        image_name = os.path.splitext(os.path.basename(image_path))[0]
 
-    for i, result in enumerate(results):
-        boxes = result.boxes
-        im_array = result.plot()  # retourne un array numpy BGR avec les boxes dessinées
+        # Sauvegarde de l'image annotée
+        result.save(filename=os.path.join(images_dir, f"{image_name}.jpg"))
 
-        plt.figure(figsize=(10, 10))
-        plt.imshow(im_array[..., ::-1])  # conversion BGR -> RGB pour matplotlib
-        plt.axis("off")
-        plt.show()
+        # Sauvegarde des prédictions au format YOLO .txt
+        save_yolo_txt(result, os.path.join(labels_dir, f"{image_name}.txt"))
+
+    print(f"Modèle testé                      : {model_name}")
+    print(f"Images annotées sauvegardées dans : {images_dir}")
+    print(f"Labels YOLO sauvegardés dans      : {labels_dir}")
 
-        result.save(filename=os.path.join(output_dir, model_dir, f"result_{i}.jpg"))  # sauvegarde avec nom unique
 
 def main() -> None:
     args = parse_args()
     predict(args.model_path, args.images_path, args.output_path)
 
+
 if __name__ == "__main__":
-    main()
- 
+    main()