update: parcellation

README.md

@@ -1,5 +1,5 @@
 ---
-title: OpenMAP T1
+title: OpenMAP-T1
 emoji: 📚
 colorFrom: gray
 colorTo: red

requirements.txt

@@ -4,9 +4,12 @@ anyio==4.9.0
 certifi==2025.1.31
 charset-normalizer==3.4.1
 click==8.1.8
+contourpy==1.3.1
+cycler==0.12.1
 fastapi==0.115.12
 ffmpy==0.5.0
 filelock==3.18.0
+fonttools==4.57.0
 fsspec==2025.3.2
 gradio==5.25.0
 gradio_client==1.8.0
@@ -17,8 +20,10 @@ httpx==0.28.1
 huggingface-hub==0.30.2
 idna==3.10
 Jinja2==3.1.6
+kiwisolver==1.4.8
 markdown-it-py==3.0.0
 MarkupSafe==3.0.2
+matplotlib==3.10.1
 mdurl==0.1.2
 mpmath==1.3.0
 networkx==3.4.2
@@ -32,6 +37,7 @@ pydantic==2.11.3
 pydantic_core==2.33.1
 pydub==0.25.1
 Pygments==2.19.1
+pyparsing==3.2.3
 python-dateutil==2.9.0.post0
 python-multipart==0.0.20
 pytz==2025.2

src/parcellation.py

@@ -1,18 +1,22 @@
 import os
-import tempfile
+import random
+import shutil
+import string
+import zipfile
 from functools import partial
 
 import gradio as gr
+import matplotlib.pyplot as plt
 import nibabel as nib
 import numpy as np
 import torch
+from PIL import Image
 from tqdm import tqdm as std_tqdm
 
 tqdm = partial(std_tqdm, dynamic_ncols=True)
 
-# 必要なモジュールのインポート
+# Import required modules from our project
 from utils.cropping import cropping
-from utils.functions import reimburse_conform
 from utils.hemisphere import hemisphere
 from utils.load_model import load_model
 from utils.make_csv import make_csv
@@ -22,46 +26,66 @@ from utils.postprocessing import postprocessing
 from utils.preprocessing import preprocessing
 from utils.stripping import stripping
 
-# モデルは起動時に一度読み込む
-MODELS = {}
 
+def nii_to_image(voxel_path, label_path, output_dir, basename):
+    """
+    Converts two NIfTI files into 2D images for visualization.
+    The voxel (input MRI) is shown as a grayscale image and the label (segmentation)
+    is shown using a default color map.
+    A middle slice is chosen by default.
+    """
+    # Load the NIfTI volumes and squeeze to remove extra dimensions
+    vdata = nib.squeeze_image(nib.as_closest_canonical(nib.load(voxel_path)))
+    ldata = nib.squeeze_image(nib.as_closest_canonical(nib.load(label_path)))
+    voxel = vdata.get_fdata().astype("float32")
+    label = ldata.get_fdata().astype("int16")
+
+    # Choose the middle slice along the first dimension and rotate for display
+    slice_index = voxel.shape[0] // 2
+    slice_voxel = np.rot90(voxel[slice_index, :, :])
+    slice_label = np.rot90(label[slice_index, :, :])
+
+    # Plot and save the input MRI image
+    plt.figure(figsize=(5, 5))
+    plt.imshow(slice_voxel, cmap="gray")
+    plt.title("Input Image")
+    plt.axis("off")
+    input_png_path = os.path.join(os.path.dirname(output_dir), f"{basename}_input.png")
+    plt.savefig(input_png_path, format="png", bbox_inches="tight", pad_inches=0)
+
+    # Plot and save the parcellation (segmentation) map image
+    plt.figure(figsize=(5, 5))
+    plt.imshow(slice_label)
+    plt.title("Parcellation Result")
+    plt.axis("off")
+    parcellation_png_path = os.path.join(
+        os.path.dirname(output_dir), f"{basename}_parcellation.png"
+    )
+    plt.savefig(parcellation_png_path, format="png", bbox_inches="tight", pad_inches=0)
 
-def load_all_models(device):
-    model_folder = "model/"  # 学習済みモデルは model/ に配置している前提
-    try:
-        # load_model 内部が argparse.Namespace 等を前提の場合、必要な属性が使われるなら適宜オブジェクトを用意してください
-        # ここでは opt を使わないように必要最低限の修正例を示しています。
-        models = load_model(model_folder=model_folder, opt=None, device=device)
-        return models  # (cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a)
-    except Exception as e:
-        print("Error during model loading:", e)
-        return None
+    return input_png_path, parcellation_png_path
 
 
 def run_inference(input_file, only_face_cropping, only_skull_stripping):
-    # 一時ディレクトリの作成
-    tmp_input_dir = tempfile.mkdtemp()
-    tmp_output_dir = tempfile.mkdtemp()
+    # Generate a random 10-character string to create a unique temporary directory
+    random_string = "".join(random.choices(string.ascii_letters + string.digits, k=10))
 
-    # アップロードされたファイルを一時フォルダに保存
+    # Extract the base filename from the uploaded file (handle .nii and .nii.gz)
     basename = os.path.splitext(os.path.basename(input_file.name))[0]
-    input_path = os.path.join(tmp_input_dir, f"{basename}.nii")
-    with open(input_path, "wb") as f:
-        f.write(input_file.read())
+    if basename.endswith(".nii"):
+        basename = os.path.splitext(basename)[0]
 
-    # Gradio用のオプションオブジェクト（オリジナルコードの argparse.Namespace 相当）
+    # Create an Options object (similar to argparse.Namespace)
     class Options:
         pass
 
     opt = Options()
-    opt.i = tmp_input_dir  # 入力ディレクトリを指定
-    opt.o = tmp_output_dir  # 出力ディレクトリを指定
-    # ここでは学習済みモデルフォルダを固定で "model/" に設定
-    opt.m = "model/"
+    # Set the output directory uniquely with the random string and base filename
+    opt.o = f"output/{random_string}/{basename}"
     opt.only_face_cropping = only_face_cropping
     opt.only_skull_stripping = only_skull_stripping
 
-    # デバイス選択
+    # Device selection: prefer CUDA if available, otherwise MPS or CPU
     if torch.cuda.is_available():
         device = torch.device("cuda")
     elif torch.backends.mps.is_available():
@@ -70,74 +94,98 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
         device = torch.device("cpu")
     print(f"Using device: {device}")
 
-    # 初回の���モデルを読み込み（キャッシュ）
-    global MODELS
-    if not MODELS:
-        MODELS = load_all_models(device)
-        if MODELS is None:
-            return "モデルの読み込みに失敗しました。"
-    cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a = MODELS
-
-    # --- 以下、元のparcellation.py の処理フローに準じた処理 ---
-    # 1. 入力画像の読み込み・キャノニカル化・squeeze
-    odata = nib.squeeze_image(nib.as_closest_canonical(nib.load(input_path)))
+    # Load the pre-trained models from the fixed "model/" folder
+    # cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
+    cnet, ssnet, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
+
+    # --- Processing Flow (based on the original parcellation.py) ---
+    # 1. Load the input image, convert to canonical orientation, and remove extra dimensions
+    odata = nib.squeeze_image(nib.as_closest_canonical(nib.load(input_file.name)))
     nii = nib.Nifti1Image(odata.get_fdata().astype(np.float32), affine=odata.affine)
-    os.makedirs(os.path.join(tmp_output_dir, "original"), exist_ok=True)
-    nib.save(nii, os.path.join(tmp_output_dir, f"original/{basename}.nii"))
+    os.makedirs(os.path.join(opt.o, "original"), exist_ok=True)
+    original_nii_path = os.path.join(opt.o, f"original/{basename}.nii")
+    nib.save(nii, original_nii_path)
 
-    # 2. 前処理
-    odata, data = preprocessing(input_path, tmp_output_dir, basename)
+    # 2. Preprocess the image
+    odata, data = preprocessing(input_file.name, opt.o, basename)
 
-    # 3. クロッピング
-    cropped = cropping(tmp_output_dir, basename, odata, data, cnet, device)
+    # 3. Cropping
+    cropped, out_filename = cropping(opt.o, basename, odata, data, cnet, device)
     if only_face_cropping:
-        return os.path.join(tmp_output_dir, f"{basename}_cropped.nii")
-
-    # 4. スキルストリッピング
-    stripped, shift = stripping(tmp_output_dir, basename, cropped, odata, data, ssnet, device)
-    if only_skull_stripping:
-        return os.path.join(tmp_output_dir, f"{basename}_stripped.nii")
-
-    # 5. パーセレーション
-    parcellated = parcellation(stripped, pnet_c, pnet_s, pnet_a, device)
-    # 6. 両半球に分離
-    separated = hemisphere(stripped, hnet_c, hnet_a, device)
-    # 7. 後処理
-    output = postprocessing(parcellated, separated, shift, device)
-    # 8. CSV作成（体積情報等）
-    df = make_csv(output, tmp_output_dir, basename)
-    # 9. パーセル結果のNIfTI作成と保存
-    nii_out = nib.Nifti1Image(output.astype(np.uint16), affine=data.affine)
-    header = odata.header
-    nii_out = nib.processing.conform(
-        nii_out,
-        out_shape=(header["dim"][1], header["dim"][2], header["dim"][3]),
-        voxel_size=(header["pixdim"][1], header["pixdim"][2], header["pixdim"][3]),
-        order=0,
+        pass
+
+    else:
+        # 4. Skull stripping
+        stripped, shift, out_filename = stripping(
+            opt.o, basename, cropped, odata, data, ssnet, device
+        )
+        if only_skull_stripping:
+            pass
+        else:
+            # 5. Parcellation
+            parcellated = parcellation(stripped, pnet_a, pnet_a, pnet_a, device)
+            # 6. Separate into hemispheres
+            separated = hemisphere(stripped, hnet_c, hnet_a, device)
+            # 7. Postprocessing
+            output = postprocessing(parcellated, separated, shift, device)
+            # 8. Create CSV with volume information, etc.
+            df = make_csv(output, opt.o, basename)
+            # 9. Create and save the parcellation result NIfTI file
+            nii_out = nib.Nifti1Image(output.astype(np.uint16), affine=data.affine)
+            header = odata.header
+            nii_out = nib.processing.conform(
+                nii_out,
+                out_shape=(header["dim"][1], header["dim"][2], header["dim"][3]),
+                voxel_size=(header["pixdim"][1], header["pixdim"][2], header["pixdim"][3]),
+                order=0,
+            )
+            out_parcellated_dir = os.path.join(opt.o, "parcellated")
+            os.makedirs(out_parcellated_dir, exist_ok=True)
+            out_filename = os.path.join(out_parcellated_dir, f"{basename}_Type1_Level5.nii")
+            nib.save(nii_out, out_filename)
+            create_parcellated_images(output, opt.o, basename, odata, data)
+
+    # Zip the entire output directory into a ZIP file
+    zip_path = os.path.join(os.path.dirname(opt.o), f"{basename}_results.zip")
+    with zipfile.ZipFile(zip_path, "w") as zipf:
+        for root, _, files in os.walk(opt.o):
+            for file in files:
+                file_path = os.path.join(root, file)
+                # Adjust the path within the zip archive
+                arcname = os.path.relpath(file_path, start=opt.o)
+                zipf.write(file_path, arcname)
+
+    # Convert the NIfTI files into visualization images (PNG)
+    input_png_path, parcellation_png_path = nii_to_image(
+        input_file.name, out_filename, opt.o, basename
     )
-    out_parcellated_dir = os.path.join(tmp_output_dir, "parcellated")
-    os.makedirs(out_parcellated_dir, exist_ok=True)
-    out_filename = os.path.join(out_parcellated_dir, f"{basename}_Type1_Level5.nii")
-    nib.save(nii_out, out_filename)
-    create_parcellated_images(output, tmp_output_dir, basename, odata, data)
 
-    # 必要に応じて、一時ファイルの削除等の処理を追加してください
+    # *** Cleanup: Remove the temporary output directory ***
+    # Note: This is performed before returning. It is not possible to execute code after the return statement.
+    shutil.rmtree(opt.o)
 
-    # 最終結果のNIfTIファイルのパスを返す
-    return out_filename
+    # Return the ZIP file path and the two visualization images
+    return zip_path, Image.open(input_png_path), Image.open(parcellation_png_path)
 
 
-# Gradioインターフェースの作成（モデルフォルダの入力は不要）
+# Create the Gradio interface (the model folder input is not needed)
 iface = gr.Interface(
     fn=run_inference,
     inputs=[
-        gr.Image(label="Input NIfTI File (.nii or .nii.gz)"),
+        gr.File(label="Input NIfTI File (.nii or .nii.gz)"),
         gr.Checkbox(label="Only Face Cropping", value=False),
         gr.Checkbox(label="Only Skull Stripping", value=False),
     ],
-    outputs=gr.Image(label="Output Parcellated NIfTI File"),
+    outputs=[
+        gr.File(label="Output Results ZIP File"),
+        gr.Image(label="MRI Image (Original)"),
+        gr.Image(label="Parcellation Map (Type1_Level5)"),
+    ],
     title="OpenMAP-T1 Inference",
-    description="学習済みモデルは model/ に配置されています。\nアップロードされたMRI画像に対してOpenMAP-T1の処理を行い、パーセル結果を返します。",
+    description=(
+        "The uploaded MRI image will be processed using OpenMAP-T1, and the parcellation "
+        "results will be returned as a ZIP file along with visualization images."
+    ),
 )
 
 if __name__ == "__main__":

src/utils/cropping.py

@@ -70,6 +70,6 @@ def cropping(output_dir, basename, odata, data, cnet, device):
     out_e = closing(out_e)
     cropped = data.get_fdata().astype("float32") * out_e
 
-    reimburse_conform(output_dir, basename, "cropped", odata, data, out_e)
+    out_filename = reimburse_conform(output_dir, basename, "cropped", odata, data, out_e)
 
-    return cropped
+    return cropped, out_filename

src/utils/functions.py

@@ -12,6 +12,7 @@ def normalize(voxel):
     voxel = (voxel * 2) - 1
     return voxel.astype("float32")
 
+
 def reimburse_conform(output_dir, basename, suffix, odata, data, output):
     nii = nib.Nifti1Image(output.astype(np.uint16), affine=data.affine)
     header = odata.header
@@ -23,8 +24,8 @@ def reimburse_conform(output_dir, basename, suffix, odata, data, output):
     )
     os.makedirs(os.path.join(output_dir, f"{suffix}"), exist_ok=True)
     nib.save(nii, os.path.join(output_dir, f"{suffix}/{basename}_{suffix}_mask.nii"))
-    
+
     result = odata.get_fdata().astype("float32") * nii.get_fdata().astype("int16")
     nii = nib.Nifti1Image(result.astype(np.float32), affine=odata.affine)
     nib.save(nii, os.path.join(output_dir, f"{suffix}/{basename}_{suffix}.nii"))
-    return
+    return os.path.join(output_dir, f"{suffix}/{basename}_{suffix}_mask.nii")

src/utils/load_model.py

@@ -4,11 +4,8 @@ import torch
 
 from utils.network import UNet
 
-CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
-PROJECT_ROOT = os.path.abspath(os.path.join(CURRENT_DIR, "..", ".."))
 
-
-def load_model(opt, device):
+def load_model(model_dir, device):
     """
     This function loads multiple pre-trained models and sets them to evaluation mode.
     The models loaded are:
@@ -27,10 +24,6 @@ def load_model(opt, device):
     Returns:
     tuple: A tuple containing all the loaded models.
     """
-    if os.path.isabs(opt.m):
-        model_dir = opt.m
-    else:
-        model_dir = os.path.join(PROJECT_ROOT, opt.m)
 
     # Load CNet model
     cnet = UNet(1, 1)
@@ -47,20 +40,20 @@ def load_model(opt, device):
     ssnet.eval()
 
     # Load PNet coronal model
-    pnet_c = UNet(3, 142)
-    pnet_c.load_state_dict(
-        torch.load(os.path.join(model_dir, "PNet", "coronal.pth"), weights_only=True)
-    )
-    pnet_c.to(device)
-    pnet_c.eval()
+    # pnet_c = UNet(3, 142)
+    # pnet_c.load_state_dict(
+    #     torch.load(os.path.join(model_dir, "PNet", "coronal.pth"), weights_only=True)
+    # )
+    # pnet_c.to(device)
+    # pnet_c.eval()
 
     # Load PNet sagittal model
-    pnet_s = UNet(3, 142)
-    pnet_s.load_state_dict(
-        torch.load(os.path.join(model_dir, "PNet", "sagittal.pth"), weights_only=True)
-    )
-    pnet_s.to(device)
-    pnet_s.eval()
+    # pnet_s = UNet(3, 142)
+    # pnet_s.load_state_dict(
+    #     torch.load(os.path.join(model_dir, "PNet", "sagittal.pth"), weights_only=True)
+    # )
+    # pnet_s.to(device)
+    # pnet_s.eval()
 
     # Load PNet axial model
     pnet_a = UNet(3, 142)
@@ -87,4 +80,5 @@ def load_model(opt, device):
     hnet_a.eval()
 
     # Return all loaded models
-    return cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a
+    # return cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a
+    return cnet, ssnet, pnet_a, hnet_c, hnet_a

src/utils/parcellation.py

@@ -74,24 +74,24 @@ def parcellation(voxel, pnet_c, pnet_s, pnet_a, device):
     sagittal = voxel
     axial = voxel.transpose(2, 1, 0)
 
-    # Perform parcellation for the coronal view
-    out_c = parcellate(coronal, pnet_c, device, "c").permute(1, 3, 0, 2)
-    torch.cuda.empty_cache()
+    # # Perform parcellation for the coronal view
+    # out_c = parcellate(coronal, pnet_c, device, "c").permute(1, 3, 0, 2)
+    # torch.cuda.empty_cache()
 
-    # Perform parcellation for the sagittal view
-    out_s = parcellate(sagittal, pnet_s, device, "s").permute(1, 0, 2, 3)
-    torch.cuda.empty_cache()
+    # # Perform parcellation for the sagittal view
+    # out_s = parcellate(sagittal, pnet_s, device, "s").permute(1, 0, 2, 3)
+    # torch.cuda.empty_cache()
 
-    # Combine the results from coronal and sagittal views
-    out_e = out_c + out_s
-    del out_c, out_s
+    # # Combine the results from coronal and sagittal views
+    # out_e = out_c + out_s
+    # del out_c, out_s
 
     # Perform parcellation for the axial view
     out_a = parcellate(axial, pnet_a, device, "a").permute(1, 3, 2, 0)
     torch.cuda.empty_cache()
 
     # Combine the results from all views
-    out_e = out_e + out_a
+    out_e = out_a  # out_e + out_a
     del out_a
 
     # Get the final parcellated output by taking the argmax

src/utils/stripping.py

@@ -82,7 +82,7 @@ def stripping(output_dir, basename, voxel, odata, data, ssnet, device):
     # Multiply the original data by the thresholded output to get the stripped brain image
     stripped = data.get_fdata().astype("float32") * out_e
 
-    reimburse_conform(output_dir, basename, "stripped", odata, data, out_e)
+    out_filename = reimburse_conform(output_dir, basename, "stripped", odata, data, out_e)
 
     # Calculate the center of mass of the stripped brain image
     x, y, z = map(int, ndimage.center_of_mass(out_e))
@@ -99,4 +99,4 @@ def stripping(output_dir, basename, voxel, odata, data, ssnet, device):
     stripped = stripped[32:-32, 16:-16, 32:-32]
 
     # Return the stripped brain image and the shifts applied
-    return stripped, (xd, yd, zd)
+    return stripped, (xd, yd, zd), out_filename