Create test_model.py

test_model.py

@@ -0,0 +1,109 @@
+import os
+import sys
+import torch
+import nibabel as nib
+import numpy as np
+from pathlib import Path
+import SimpleITK as sitk
+
+# =======================
+# CONFIG
+# =======================
+SAVE_DIR = "./results"
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+MODEL_PATH = "model.best"   # ton modèle entraîné
+
+# =======================
+# CHARGEMENT DU MODÈLE
+# =======================
+def load_model():
+    print("🔄 Chargement du modèle...")
+    model = torch.load(MODEL_PATH, map_location="cpu")
+    model.eval()
+    return model
+
+# =======================
+# PRÉPROCESSING
+# =======================
+def load_image(file_path):
+    ext = Path(file_path).suffix.lower()
+    
+    if ext in [".nii", ".gz"]:
+        img = nib.load(file_path)
+        data = img.get_fdata()
+        affine = img.affine
+        return data, affine
+    
+    elif ext == ".dcm":
+        itk_img = sitk.ReadImage(file_path)
+        data = sitk.GetArrayFromImage(itk_img)
+        affine = np.eye(4)  # placeholder simple
+        return data, affine
+    
+    else:
+        raise ValueError("❌ Format non supporté (seulement NIfTI ou DICOM).")
+
+# =======================
+# INFERENCE
+# =======================
+def run_inference(model, data):
+    data = (data - np.min(data)) / (np.max(data) - np.min(data) + 1e-8)
+    data_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # (B,C,H,W,D)
+    
+    with torch.no_grad():
+        pred = model(data_tensor)
+        pred = torch.argmax(pred, dim=1).squeeze().cpu().numpy()
+    return pred
+
+# =======================
+# RAPPORT
+# =======================
+def generate_report(pred, save_path):
+    volume_total = np.prod(pred.shape)
+    volume_tumeur = np.sum(pred > 0)
+    ratio = (volume_tumeur / volume_total) * 100
+    
+    report = [
+        "📄 RAPPORT MÉDICAL AUTOMATISÉ",
+        "-----------------------------------",
+        f"Dimensions de l'image : {pred.shape}",
+        f"Volume total voxels   : {volume_total}",
+        f"Volume suspect tumeur : {volume_tumeur}",
+        f"Ratio suspect         : {ratio:.2f} %",
+        "",
+        "⚠️ Rapport automatique à valider par un médecin."
+    ]
+    
+    with open(save_path, "w", encoding="utf-8") as f:
+        f.write("\n".join(report))
+    
+    print("\n".join(report))
+
+# =======================
+# MAIN
+# =======================
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("❌ Utilisation : python test_model.py <chemin_fichier_IRM>")
+        sys.exit(1)
+    
+    input_path = sys.argv[1]
+    print(f"📂 Fichier en entrée : {input_path}")
+
+    # Load
+    model = load_model()
+    data, affine = load_image(input_path)
+
+    # Inference
+    pred = run_inference(model, data)
+
+    # Save segmentation
+    out_seg_path = os.path.join(SAVE_DIR, "segmentation_pred.nii.gz")
+    nib.save(nib.Nifti1Image(pred.astype(np.uint8), affine), out_seg_path)
+    print(f"✅ Segmentation sauvegardée : {out_seg_path}")
+
+    # Save report
+    out_report_path = os.path.join(SAVE_DIR, "segmentation_report_detailed.txt")
+    generate_report(pred, out_report_path)
+    print(f"✅ Rapport sauvegardé : {out_report_path}")