File size: 3,497 Bytes
f0064c6 fe65029 f0064c6 d29468e ef8cc33 f0064c6 ef8cc33 f0064c6 ef8cc33 f0064c6 fe65029 f0064c6 ef8cc33 fe65029 f0064c6 fe65029 f0064c6 fe65029 ef8cc33 f0064c6 fe65029 ef8cc33 f0064c6 fe65029 ef8cc33 f0064c6 ef8cc33 fe65029 ef8cc33 f0064c6 fe65029 f0064c6 ef8cc33 f0064c6 ef8cc33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import argparse
import glob
import os
import torch
from ultralytics import YOLO
# uv run inference.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 et des labels .txt"
)
parser.add_argument(
"model_path",
type=str,
help="Chemin vers le modèle YOLO (.pt)",
)
parser.add_argument(
"images_path",
type=str,
help="Dossier contenant les images à traiter",
)
parser.add_argument(
"output_path",
type=str,
help="Dossier de sortie pour les images annotées et les labels",
)
parser.add_argument(
"--imgsz",
type=int,
default=640,
help="Taille d'image pour l'inférence (défaut: 640)",
)
parser.add_argument(
"--half",
action="store_true",
help="Utiliser la précision FP16 pour réduire la consommation mémoire",
)
return parser.parse_args()
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, imgsz: int, half: bool):
model = YOLO(model_path)
images = glob.glob(os.path.join(images_path, "*.jpg"))
selection = images
print(f"{len(selection)} images sélectionnées pour la prédiction")
# 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)
total = len(selection)
for i, image_path in enumerate(selection, start=1):
image_name = os.path.splitext(os.path.basename(image_path))[0]
# Inférence sur une seule image à la fois
with torch.no_grad():
results = model(image_path, imgsz=imgsz, verbose=False)
result = results[0]
# Sauvegarde de l'image annotée
result.save(filename=os.path.join(images_dir, f"{image_name}.jpg"))
# Sauvegarde des prédictions au format YOLO .txt
save_yolo_txt(result, os.path.join(labels_dir, f"{image_name}.txt"))
# Libération explicite de la mémoire GPU après chaque image
del results, result
torch.cuda.empty_cache()
if i % 10 == 0 or i == total:
print(f"[{i}/{total}] traitées")
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}")
def main() -> None:
args = parse_args()
predict(args.model_path, args.images_path, args.output_path, args.imgsz, args.half)
if __name__ == "__main__":
main() |