3d->4d proj

README.md

@@ -10,5 +10,3 @@ pinned: false
 license: mit
 short_description: 'fMRI Learning Stage Classification with Vision Transformers '
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -14,8 +14,9 @@ from einops.layers.torch import Rearrange
 from scipy.ndimage import zoom
 import matplotlib.pyplot as plt
 import seaborn as sns
+from nilearn import plotting
+import matplotlib.gridspec as gridspec
 
-# core config
 @dataclass
 class Config:
     VOLUME_SIZE: Tuple[int, int, int] = (64, 64, 30)
@@ -25,7 +26,6 @@ class Config:
     DROPOUT: float = 0.1
     TASK_DIM: int = 512
 
-# model components
 class HierarchicalAttention(nn.Module):
     def __init__(self, dim, heads=8):
         super().__init__()
@@ -206,19 +206,55 @@ def preprocess_volume(vol, target_size=(64, 64, 30)):
     vol = (vol - vol.mean((1,2,3,4), keepdims=True)) / (vol.std((1,2,3,4), keepdims=True) + 1e-8)
     return torch.from_numpy(vol).float()
 
-def plot_results(region_acts, temporal_pattern):
-    fig = plt.figure(figsize=(12,4))
+def plot_brain_slices(data, learning_stage):
+    fig = plt.figure(figsize=(15, 5))
+    mean_activation = data.mean(axis=0)
     
-    plt.subplot(121)
-    sns.heatmap(region_acts.reshape(1,-1), cmap='RdBu_r', center=0)
-    plt.title('region activations')
-    plt.xlabel('brain region')
+    for i, slice_idx in enumerate([mean_activation.shape[-1]//4, 
+                                 mean_activation.shape[-1]//2, 
+                                 3*mean_activation.shape[-1]//4]):
+        plt.subplot(1, 3, i+1)
+        plt.imshow(mean_activation[...,slice_idx].T, cmap='hot')
+        plt.colorbar()
+        plt.title(f'slice z={slice_idx}\nlearning: {learning_stage:.3f}')
+        plt.axis('off')
     
-    plt.subplot(122)
-    plt.plot(temporal_pattern.squeeze())
-    plt.title('temporal pattern')
-    plt.xlabel('time')
+    return fig
+
+def plot_results(data, region_acts, temporal_pattern, learning_stage):
+    fig = plt.figure(figsize=(15,10))
+    gs = gridspec.GridSpec(2, 2)
+    
+    # brain slices
+    ax1 = plt.subplot(gs[0,:])
+    mean_activation = data.mean(axis=0)
+    slice_idx = mean_activation.shape[-1]//2
+    brain_slice = mean_activation[...,slice_idx]
+    
+    # find most active region
+    peak_coords = np.unravel_index(np.argmax(brain_slice), brain_slice.shape)
+    
+    im = ax1.imshow(brain_slice.T, cmap='hot')
+    plt.colorbar(im, ax=ax1)
+    ax1.plot(peak_coords[0], peak_coords[1], 'r*', markersize=15, 
+             label=f'peak ({peak_coords[0]}, {peak_coords[1]})')
+    ax1.legend()
+    ax1.set_title(f'brain activation (z={slice_idx})\nlearning stage: {learning_stage:.3f}')
+    
+    # region activations
+    ax2 = plt.subplot(gs[1,0])
+    max_region = np.argmax(region_acts)
+    sns.heatmap(region_acts.reshape(1,-1), cmap='RdBu_r', center=0, ax=ax2)
+    ax2.set_title(f'region activations\nmost active: {max_region}')
+    ax2.set_xlabel('brain region')
+    
+    # temporal pattern
+    ax3 = plt.subplot(gs[1,1])
+    ax3.plot(temporal_pattern.squeeze())
+    ax3.set_title('temporal dynamics')
+    ax3.set_xlabel('time')
     
+    plt.tight_layout()
     return fig
 
 def process_fmri(file_obj):
@@ -252,8 +288,11 @@ def process_fmri(file_obj):
         for stage in ['full', 'region', 'temporal']:
             try:
                 model = SequentialBrainViT(Config())
+                model._init_weights()  # critical: init before load
                 ckpt = torch.load(f'best_{stage}.pt', map_location=device)
-                model.load_state_dict(ckpt['model'])
+                missing = model.load_state_dict(ckpt['model'], strict=False)
+                if missing:
+                    print(f"warning - {stage} missing keys:", missing)
                 model.eval()
                 
                 with torch.no_grad():
@@ -265,8 +304,10 @@ def process_fmri(file_obj):
                     }
                     
                     fig = plot_results(
+                        data[0].cpu().numpy(),  # drop batch
                         results[stage]['region_activation'],
-                        results[stage]['temporal_pattern']
+                        results[stage]['temporal_pattern'],
+                        results[stage]['learning_stage']
                     )
                     figs.append(fig)
                     plt.close()
@@ -274,9 +315,12 @@ def process_fmri(file_obj):
             except Exception as e:
                 return f"error in {stage} model: {str(e)}", None
         
+        # enhanced results text w/ peak info
         stage_results = "\n".join([
             f"{stage.upper()} MODEL:"
             f"\nlearning stage: {res['learning_stage']:.3f}"
+            f"\npeak region: {np.argmax(res['region_activation'])}"
+            f"\npeak activation: {np.max(res['region_activation']):.3f}"
             f"\n"
             for stage, res in results.items()
         ])
@@ -286,13 +330,12 @@ def process_fmri(file_obj):
     except Exception as e:
         return f"error processing file: {str(e)}", None
 
-# create interface
 iface = gr.Interface(
     fn=process_fmri,
     inputs=gr.File(label="upload 4D fMRI nifti (.nii/.nii.gz)"),
     outputs=[
         gr.Textbox(label="classification results"),
-        gr.Plot(label="visualization")
+        gr.Plot(label="brain activation + analysis")
     ],
     title="fmri learning stage classifier",
     description="upload a 4D fMRI nifti file to classify learning stages and visualize brain patterns",