shape range val

app.py

@@ -226,14 +226,23 @@ def process_fmri(file_obj):
         img = nib.load(file_obj.name)
         data = img.get_fdata(dtype=np.float32)
         
-        if data.ndim != 4:
-            return f"error: expected 4D data, got {data.ndim}D", None
-        
+        # shape validation + expansion
         if data.ndim == 3:
-            data = data[None,...]  # add time dim [H,W,D] -> [1,H,W,D]
+            data = data[None,...]  # [H,W,D] -> [T,H,W,D]
         elif data.ndim != 4:
-            return f"error: expected 3D/4D data, got {data.ndim}D", None
+            return f"error: volume must be 3D/4D, got {data.ndim}D", None
             
+        # validate dims
+        t,h,w,d = data.shape
+        if t < 1 or h < 16 or w < 16 or d < 8:
+            return f"error: invalid dims {data.shape}, min: [1,16,16,8]", None
+        if t > 1000 or h > 256 or w > 256 or d > 256:
+            return f"error: dims too large {data.shape}, max: [1000,256,256,256]", None
+        
+        # reshape for batch
+        data = data.reshape(1, t, h, w, d)  # [B,T,H,W,D]
+        
+        # normalize + preprocess
         data = preprocess_volume(data)
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
@@ -241,25 +250,29 @@ def process_fmri(file_obj):
         figs = []
         
         for stage in ['full', 'region', 'temporal']:
-            model = SequentialBrainViT(Config())
-            ckpt = torch.load(f'best_{stage}.pt', map_location=device)
-            model.load_state_dict(ckpt['model'])
-            model.eval()
-            
-            with torch.no_grad():
-                outputs = model(data.to(device), torch.tensor([0]).to(device))
-                results[stage] = {
-                    'learning_stage': float(outputs['learning_stage'].cpu().mean()),
-                    'region_activation': outputs['region_activation'].cpu().numpy(),
-                    'temporal_pattern': outputs['temporal_pattern'].cpu().numpy()
-                }
+            try:
+                model = SequentialBrainViT(Config())
+                ckpt = torch.load(f'best_{stage}.pt', map_location=device)
+                model.load_state_dict(ckpt['model'])
+                model.eval()
                 
-                fig = plot_results(
-                    results[stage]['region_activation'],
-                    results[stage]['temporal_pattern']
-                )
-                figs.append(fig)
-                plt.close()
+                with torch.no_grad():
+                    outputs = model(data.to(device), torch.tensor([0]).to(device))
+                    results[stage] = {
+                        'learning_stage': float(outputs['learning_stage'].cpu().mean()),
+                        'region_activation': outputs['region_activation'].cpu().numpy(),
+                        'temporal_pattern': outputs['temporal_pattern'].cpu().numpy()
+                    }
+                    
+                    fig = plot_results(
+                        results[stage]['region_activation'],
+                        results[stage]['temporal_pattern']
+                    )
+                    figs.append(fig)
+                    plt.close()
+                    
+            except Exception as e:
+                return f"error in {stage} model: {str(e)}", None
         
         stage_results = "\n".join([
             f"{stage.upper()} MODEL:"
@@ -268,7 +281,7 @@ def process_fmri(file_obj):
             for stage, res in results.items()
         ])
         
-        return stage_results, figs[0]  # return first fig for display
+        return stage_results, figs[0]
         
     except Exception as e:
         return f"error processing file: {str(e)}", None