Upload 16 files

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.nii filter=lfs diff=lfs merge=lfs -text
+*.nii.gz filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,9 @@
+Apache License 2.0
+
+Copyright 2025 hmnshudhmn24
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0

README.md

@@ -1,3 +1,58 @@
----
-license: apache-2.0
----
+---
+language:
+  - en
+license: apache-2.0
+tags:
+  - image-to-3d
+  - medical-imaging
+  - reconstruction
+  - segmentation
+  - explainability
+pipeline_tag: image-to-3d
+library_name: python
+---
+
+# MEDIVIEW-3D (Advanced)
+
+**MEDIVIEW-3D** converts 2D medical image slices into a 3D reconstruction, localizes anomalous regions, and generates **textual explanations** describing the detected regions (size, approximate location, and suggested next steps).
+
+**Important:** This is a research/demo tool and **not** a medical device. Do not use for clinical decisions.
+
+## Quickstart (demo)
+
+1. Install:
+```bash
+pip install -r requirements.txt
+```
+
+2. Generate synthetic phantom slices:
+```bash
+python examples/generate_synthetic_phantom.py
+```
+
+3. Run inference with thresholding and get a mesh + explanation:
+```bash
+python infer_anomaly.py --source examples/synthetic_phantom --method threshold --out demo_mesh_threshold.ply --explain_out explanation.txt
+```
+
+4. (Optional) Train small UNet and run model-based inference:
+```bash
+python train_unet.py --data examples/synthetic_phantom --epochs 3 --out models/unet_demo.pt
+python infer_anomaly.py --source examples/synthetic_phantom --method model --model_path models/unet_demo.pt --out demo_mesh_model.ply --explain_out explanation_model.txt
+```
+
+5. Run Streamlit demo:
+```bash
+streamlit run app.py
+```
+
+## What you get
+- 3D mesh `.ply` with anomaly regions colored red
+- `explanation.txt` with human-friendly descriptions of detected regions
+- Example synthetic phantom (no patient data)
+- Small UNet implementation for demo training
+
+## Safety & Limitations
+- Demo-only; not clinically validated.
+- Do not upload identifiable patient data to public repos.
+- For real medical use, integrate robust preprocessing and obtain regulatory approvals.

app.py

@@ -0,0 +1,52 @@
+import streamlit as st
+from PIL import Image
+import numpy as np
+import tempfile, os, zipfile
+from io_utils import load_slices_from_folder
+from infer_anomaly import simple_threshold_mask, mesh_and_mask_from_volume
+from io_utils import save_mesh_ply
+
+st.set_page_config(page_title="MEDIVIEW-3D", layout="wide")
+st.title("MEDIVIEW-3D — 2D → 3D Reconstruction + Explanation")
+
+uploaded = st.file_uploader("Upload a zip of PNG slices (or skip to use example)", type=["zip"])
+use_example = st.button("Use example synthetic slices")
+thresh = st.slider("Threshold (for simple anomaly)", 0.0, 1.0, 0.6)
+out_mesh = st.text_input("Output mesh path", value="mesh_demo.ply")
+
+def unpack_and_read(zipf):
+    tmpdir = tempfile.mkdtemp()
+    with zipfile.ZipFile(zipf) as z:
+        z.extractall(tmpdir)
+    vol = load_slices_from_folder(tmpdir, glob_pattern="*.png")
+    return vol, tmpdir
+
+vol = None
+if uploaded:
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp:
+        tmp.write(uploaded.getvalue())
+        tmp.flush()
+        vol, tmpdir = unpack_and_read(tmp.name)
+elif use_example:
+    if os.path.isdir("examples/synthetic_phantom"):
+        vol = load_slices_from_folder("examples/synthetic_phantom", glob_pattern="*.png")
+    else:
+        st.error("Example slices not found in examples/synthetic_phantom")
+
+if vol is not None:
+    st.write("Volume shape (z,y,x):", vol.shape)
+    mid = vol.shape[0]//2
+    st.image(Image.fromarray(vol[mid].astype('uint8')), caption=f"Middle slice (index {mid})")
+    if st.button("Run reconstruction & anomaly detection"):
+        mask = simple_threshold_mask(vol, threshold=thresh)
+        verts, faces, normals, colors = mesh_and_mask_from_volume(vol, iso=0.5, spacing=(1.0,1.0,1.0), mask=mask)
+        save_mesh_ply(verts, faces, out_mesh, normals=normals, colors=colors)
+        st.success(f"Saved colored mesh to {out_mesh}")
+        # Generate simple textual explanation
+        from infer_anomaly import analyze_regions, generate_text_explanation
+        regions = analyze_regions(mask, min_voxels=20)
+        expl = generate_text_explanation(regions, vol.shape)
+        st.subheader('Detected regions explanation')
+        st.text(expl)
+else:
+    st.info("Upload a slice zip or click 'Use example synthetic slices' to proceed.")

config.yaml

@@ -0,0 +1,13 @@
+io:
+  slice_glob: "examples/synthetic_phantom/*.png"
+reconstruct:
+  iso_value: 0.5
+  spacing: [1.0, 1.0, 1.0]
+anomaly:
+  method: "threshold"
+  threshold: 0.6
+model:
+  path: "models/unet_best.pt"
+text_explainer:
+  enabled: true
+  min_region_voxels: 50

examples/demo_commands.txt

@@ -0,0 +1,13 @@
+# Demo commands
+
+# 1) Generate synthetic phantom slices
+python examples/generate_synthetic_phantom.py
+
+# 2) Reconstruct mesh + explanation (thresholding)
+python infer_anomaly.py --source examples/synthetic_phantom --method threshold --out demo_mesh_threshold.ply --explain_out explanation.txt
+
+# 3) Train small UNet (optional)
+python train_unet.py --data examples/synthetic_phantom --epochs 3 --out models/unet_demo.pt
+
+# 4) Run model-based inference (after training)
+python infer_anomaly.py --source examples/synthetic_phantom --method model --model_path models/unet_demo.pt --out demo_mesh_model.ply --explain_out explanation_model.txt

examples/generate_synthetic_phantom.py

@@ -0,0 +1,15 @@
+import os, numpy as np
+from PIL import Image
+os.makedirs("examples/synthetic_phantom", exist_ok=True)
+z = 64
+y = 128
+x = 128
+for i in range(z):
+    Y, X = np.meshgrid(np.linspace(-1,1,x), np.linspace(-1,1,y))
+    cx = 0.2*np.sin(2*np.pi*i/z)
+    cy = 0.2*np.cos(2*np.pi*i/z)
+    r = np.sqrt((X-cx)**2 + (Y-cy)**2)
+    img = (np.exp(-(r**2)/(2*(0.2**2))) * 255).astype(np.uint8)
+    img = img + np.random.randint(0,20,size=img.shape).astype(np.uint8)
+    Image.fromarray(img).save(f"examples/synthetic_phantom/slice_{i:03d}.png")
+print("Synthetic phantom saved to examples/synthetic_phantom/")

infer_anomaly.py

@@ -0,0 +1,129 @@
+import argparse, yaml, numpy as np
+from utils import normalize_volume
+from io_utils import load_slices_from_folder, load_nifti, load_dicom_folder, save_mesh_ply
+from skimage import measure
+import os
+
+def simple_threshold_mask(volume, threshold=0.6):
+    vol_n = normalize_volume(volume)
+    return (vol_n > threshold).astype(np.uint8)
+
+def mask_to_vertex_colors(verts, faces, mask_vol, spacing=(1.0,1.0,1.0)):
+    coords = verts / np.array(spacing)
+    coords_idx = np.round(coords).astype(int)
+    colors = np.zeros((len(verts), 3), dtype=float)
+    zmax, ymax, xmax = mask_vol.shape
+    for i, (z,y,x) in enumerate(coords_idx):
+        if 0 <= z < zmax and 0 <= y < ymax and 0 <= x < xmax:
+            if mask_vol[z,y,x]:
+                colors[i] = np.array([1.0, 0.0, 0.0])
+            else:
+                colors[i] = np.array([0.7, 0.7, 0.7])
+        else:
+            colors[i] = np.array([0.7, 0.7, 0.7])
+    return colors
+
+def mesh_and_mask_from_volume(vol, iso=0.5, spacing=(1.0,1.0,1.0), mask=None):
+    vol_n = normalize_volume(vol)
+    verts, faces, normals, values = measure.marching_cubes(vol_n, level=iso, spacing=spacing)
+    colors = None
+    if mask is not None:
+        colors = mask_to_vertex_colors(verts, faces, mask, spacing=spacing)
+    return verts, faces, normals, colors
+
+def analyze_regions(mask, min_voxels=50):
+    # find connected components and simple stats
+    from scipy import ndimage as ndi
+    labeled, n = ndi.label(mask)
+    regions = []
+    for lab in range(1, n+1):
+        coords = np.argwhere(labeled==lab)
+        voxels = coords.shape[0]
+        z_mean, y_mean, x_mean = coords.mean(axis=0).tolist()
+        regions.append({'label': lab, 'voxels': int(voxels), 'center': [float(z_mean), float(y_mean), float(x_mean)]})
+    # filter small
+    regions = [r for r in regions if r['voxels'] >= min_voxels]
+    return regions
+
+def generate_text_explanation(regions, vol_shape, spacing=(1.0,1.0,1.0)):
+    if not regions:
+        return "No anomalous regions detected above threshold."
+
+    texts = []
+
+    zdim, ydim, xdim = vol_shape
+    for r in regions:
+        zc, yc, xc = r['center']
+        # approximate location as top/middle/bottom and left/center/right
+        zpos = 'top' if zc < zdim*0.33 else ('bottom' if zc > zdim*0.66 else 'middle')
+        ypos = 'left' if xc < xdim*0.33 else ('right' if xc > xdim*0.66 else 'center')
+        texts.append(f"Region {r['label']}: approx {r['voxels']} voxels, located near the {zpos} (z~{zc:.1f}), {ypos} (x~{xc:.1f}). Suggest clinical review and consider high-resolution imaging or segmentation.")
+    return "\n".join(texts)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", required=True)
+    parser.add_argument("--source_type", choices=['folder','dicom','nifti'], default='folder')
+    parser.add_argument("--glob", default="*.png")
+    parser.add_argument("--config", default="config.yaml")
+    parser.add_argument("--method", choices=['threshold','model'], default='threshold')
+    parser.add_argument("--threshold", type=float, default=None)
+    parser.add_argument("--model_path", default=None)
+    parser.add_argument("--out", default="mesh_colored.ply")
+    parser.add_argument("--explain_out", default="explanation.txt")
+    args = parser.parse_args()
+
+    try:
+        cfg = yaml.safe_load(open(args.config))
+    except Exception:
+        cfg = {}
+    cfg_thresh = cfg.get('anomaly', {}).get('threshold', 0.6)
+    threshold = args.threshold or cfg_thresh
+    spacing = tuple(cfg.get('reconstruct', {}).get('spacing', [1.0,1.0,1.0]))
+    iso = cfg.get('reconstruct',{}).get('iso_value', 0.5)
+    min_vox = cfg.get('text_explainer',{}).get('min_region_voxels', 50)
+
+    # load volume
+    if args.source_type == 'folder':
+        vol = load_slices_from_folder(args.source, glob_pattern=args.glob)
+    elif args.source_type == 'dicom':
+        vol = load_dicom_folder(args.source)
+    else:
+        vol = load_nifti(args.source)
+
+    if args.method == 'threshold':
+        mask = simple_threshold_mask(vol, threshold=threshold)
+    else:
+        # model-based per-slice segmentation (if model provided)
+        try:
+            import torch
+            from models.unet import UNet
+            model = UNet(in_channels=1, out_channels=1)
+            model.load_state_dict(torch.load(args.model_path, map_location='cpu'))
+            model.eval()
+            vol_n = normalize_volume(vol)
+            mask = np.zeros_like(vol_n, dtype=np.uint8)
+            for i in range(vol_n.shape[0]):
+                s = vol_n[i]
+                x = (s - s.min())/(s.max()-s.min()+1e-8)
+                import torch
+                inp = torch.tensor(x[np.newaxis, np.newaxis, ...], dtype=torch.float32)
+                with torch.no_grad():
+                    out = model(inp).numpy()[0,0]
+                mask[i] = (out > 0.5).astype(np.uint8)
+        except Exception as e:
+            print("Model-based method failed:", e)
+            mask = simple_threshold_mask(vol, threshold=threshold)
+
+    regions = analyze_regions(mask, min_voxels=min_vox)
+    explanation = generate_text_explanation(regions, vol.shape, spacing=spacing)
+    verts, faces, normals, colors = mesh_and_mask_from_volume(vol, iso=iso, spacing=spacing, mask=mask)
+    save_mesh_ply(verts, faces, args.out, normals=normals, colors=colors)
+    with open(args.explain_out, 'w') as f:
+        f.write(explanation)
+    print("Saved colored mesh to", args.out)
+    print("Saved textual explanation to", args.explain_out)
+
+if __name__ == '__main__':
+    main()

io_utils.py

@@ -0,0 +1,62 @@
+import os, glob
+import numpy as np
+from PIL import Image
+import pydicom
+import nibabel as nib
+
+def load_slices_from_folder(folder, glob_pattern="*.png", sort_numerically=True):
+    paths = sorted(glob.glob(os.path.join(folder, glob_pattern)))
+    if sort_numerically:
+        try:
+            paths = sorted(paths, key=lambda p: int(''.join(filter(str.isdigit, os.path.basename(p))) or 0))
+        except Exception:
+            pass
+    slices = [np.array(Image.open(p).convert("L")) for p in paths]
+    vol = np.stack(slices, axis=0)  # z, y, x
+    return vol.astype(np.float32)
+
+def load_dicom_folder(folder):
+    files = [p for p in glob.glob(os.path.join(folder, "**"), recursive=False) if os.path.isfile(p)]
+    ds = []
+    for f in files:
+        try:
+            d = pydicom.dcmread(f)
+            ds.append((getattr(d, "InstanceNumber", 0), d))
+        except Exception:
+            continue
+    ds = sorted(ds, key=lambda x: x[0])
+    imgs = [d.pixel_array for _, d in ds]
+    vol = np.stack(imgs, axis=0).astype(np.float32)
+    return vol
+
+def load_nifti(path):
+    nii = nib.load(path)
+    arr = nii.get_fdata().astype(np.float32)
+    if arr.ndim == 3:
+        return np.transpose(arr, (2,1,0))
+    return arr
+
+def save_mesh_ply(verts, faces, out_path, normals=None, colors=None):
+    import numpy as np
+    with open(out_path, 'w') as f:
+        f.write("ply\nformat ascii 1.0\n")
+        f.write(f"element vertex {len(verts)}\n")
+        f.write("property float x\nproperty float y\nproperty float z\n")
+        if colors is not None:
+            f.write("property uchar red\nproperty uchar green\nproperty uchar blue\n")
+        if normals is not None:
+            f.write("property float nx\nproperty float ny\nproperty nz\n")
+        f.write(f"element face {len(faces)}\n")
+        f.write("property list uchar int vertex_indices\n")
+        f.write("end_header\n")
+        for i, v in enumerate(verts):
+            line = f"{v[0]} {v[1]} {v[2]}"
+            if colors is not None:
+                c = (np.clip(colors[i], 0, 1) * 255).astype(int)
+                line += f" {c[0]} {c[1]} {c[2]}"
+            if normals is not None:
+                n = normals[i]
+                line += f" {n[0]} {n[1]} {n[2]}"
+            f.write(line + "\n")
+        for face in faces:
+            f.write(f"3 {face[0]} {face[1]} {face[2]}\n")

models/unet.py

@@ -0,0 +1,54 @@
+import torch
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, 3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+    def forward(self, x): return self.net(x)
+
+class Down(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.net = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch))
+    def forward(self, x): return self.net(x)
+
+class Up(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(in_ch, in_ch//2, 2, stride=2)
+        self.conv = DoubleConv(in_ch, out_ch)
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+        x1 = F.pad(x1, [diffX//2, diffX - diffX//2, diffY//2, diffY - diffY//2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, base_c=32):
+        super().__init__()
+        self.inc = DoubleConv(in_channels, base_c)
+        self.down1 = Down(base_c, base_c*2)
+        self.down2 = Down(base_c*2, base_c*4)
+        self.up1 = Up(base_c*4, base_c*2)
+        self.up2 = Up(base_c*2, base_c)
+        self.outc = nn.Conv2d(base_c, out_channels, 1)
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x = self.up1(x3, x2)
+        x = self.up2(x, x1)
+        x = self.outc(x)
+        return torch.sigmoid(x)

notebooks/MEDIVIEW-3D_Demo.ipynb

@@ -0,0 +1 @@
+{"cells": [{"cell_type": "markdown", "metadata": {}, "source": ["# MEDIVIEW-3D Demo Notebook\\n", "Run the demo scripts locally."]}, {"cell_type": "code", "metadata": {}, "source": ["!python examples/generate_synthetic_phantom.py\\n", "!python infer_anomaly.py --source examples/synthetic_phantom --method threshold --out demo_mesh_threshold.ply --explain_out explanation.txt\\n"]}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python"}}, "nbformat": 4, "nbformat_minor": 5}

reconstruct_3d.py

@@ -0,0 +1,45 @@
+import argparse, yaml
+import numpy as np
+from skimage import measure
+from utils import normalize_volume
+from io_utils import load_slices_from_folder, load_dicom_folder, load_nifti, save_mesh_ply
+
+def build_volume(source, source_type='folder', glob_pattern="*.png"):
+    if source_type == 'folder':
+        vol = load_slices_from_folder(source, glob_pattern=glob_pattern)
+    elif source_type == 'dicom':
+        vol = load_dicom_folder(source)
+    elif source_type == 'nifti':
+        vol = load_nifti(source)
+    else:
+        raise ValueError("Unsupported source_type")
+    return vol
+
+def mesh_from_volume(vol, iso=0.5, spacing=(1.0,1.0,1.0)):
+    vol_n = normalize_volume(vol)
+    verts, faces, normals, values = measure.marching_cubes(vol_n, level=iso, spacing=spacing)
+    return verts, faces, normals
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", required=True)
+    parser.add_argument("--source_type", choices=['folder','dicom','nifti'], default='folder')
+    parser.add_argument("--glob", default="*.png")
+    parser.add_argument("--config", default="config.yaml")
+    parser.add_argument("--out", default="mesh.ply")
+    parser.add_argument("--iso", type=float, default=None)
+    args = parser.parse_args()
+
+    try:
+        cfg = yaml.safe_load(open(args.config))
+    except Exception:
+        cfg = {}
+    iso = args.iso or (cfg.get('reconstruct',{}).get('iso_value', 0.5))
+    spacing = cfg.get('reconstruct',{}).get('spacing', [1.0,1.0,1.0])
+    vol = build_volume(args.source, source_type=args.source_type, glob_pattern=args.glob)
+    verts, faces, normals = mesh_from_volume(vol, iso=iso, spacing=tuple(spacing))
+    save_mesh_ply(verts, faces, args.out, normals=normals)
+    print(f"Saved mesh to {args.out} (verts={len(verts)}, faces={len(faces)})")
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,13 @@
+numpy>=1.21
+scipy
+scikit-image
+matplotlib
+pillow
+imageio
+pydicom
+nibabel
+PyYAML
+trimesh
+torch>=1.10
+tqdm
+streamlit

train_unet.py

@@ -0,0 +1,48 @@
+import os, glob
+import numpy as np
+from PIL import Image
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torch.nn as nn
+from models.unet import UNet
+import torch.optim as optim
+
+class SliceDataset(Dataset):
+    def __init__(self, folder):
+        self.paths = sorted(glob.glob(os.path.join(folder, "*.png")))
+    def __len__(self): return len(self.paths)
+    def __getitem__(self, idx):
+        p = self.paths[idx]
+        img = np.array(Image.open(p).convert("L"), dtype=np.float32)/255.0
+        mask = (img > img.mean() + 0.25).astype(np.float32)
+        img = img[np.newaxis,...]
+        mask = mask[np.newaxis,...]
+        return torch.tensor(img), torch.tensor(mask)
+
+def train(folder, epochs=3, out="models/unet_best.pt"):
+    ds = SliceDataset(folder)
+    dl = DataLoader(ds, batch_size=4, shuffle=True)
+    model = UNet(in_channels=1, out_channels=1)
+    opt = optim.Adam(model.parameters(), lr=1e-3)
+    loss_fn = nn.BCELoss()
+    for epoch in range(epochs):
+        total=0
+        model.train()
+        for x,y in dl:
+            outp = model(x)
+            loss = loss_fn(outp, y)
+            opt.zero_grad(); loss.backward(); opt.step()
+            total += loss.item()
+        print(f"Epoch {epoch+1}, loss {total/len(dl):.4f}")
+    os.makedirs(os.path.dirname(out), exist_ok=True)
+    torch.save(model.state_dict(), out)
+    print("Saved model to", out)
+
+if __name__ == '__main__':
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument('--data', default='examples/synthetic_phantom')
+    p.add_argument('--epochs', type=int, default=3)
+    p.add_argument('--out', default='models/unet_best.pt')
+    args = p.parse_args()
+    train(args.data, epochs=args.epochs, out=args.out)

utils.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+def normalize_volume(vol):
+    vmin = np.nanmin(vol)
+    vmax = np.nanmax(vol)
+    if vmax - vmin < 1e-8:
+        return np.zeros_like(vol)
+    return (vol - vmin) / (vmax - vmin)
+
+def ensure_3d(volume):
+    if volume.ndim == 2:
+        return volume[np.newaxis, ...]
+    return volume