Hugging Face's logo

Spaces:
Stroke-ia
/
STREAMLITE
Sleeping

App Files Community
STREAMLITE
File size: 3,144 Bytes 
6f78e96
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
109
110
 
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}")