added small changes

Author: Ian-Goodall-Halliwell

micaflow/scripts/compute_fa_md.py

@@ -472,6 +472,28 @@ def compute_fa_md(bias_corr_path, mask_path, moving_bval, moving_bvec, fa_path,
         print(f"  New shape: {dwi_data.shape}")
         print(f"  Total volumes: {len(bvals)}")
 
+    # Filter out invalid diffusion volumes (b > 50 but bvec is 0,0,0)
+    # bvecs is shaped (3, N) at this point
+    bvec_norms = np.linalg.norm(bvecs, axis=0)
+    # Identify indices that are in a shell (b>50) but have no direction
+    bad_indices = (bvals > 50) & (bvec_norms < 1e-6)
+    
+    if np.any(bad_indices):
+        n_removed = np.sum(bad_indices)
+        print(f"{YELLOW}Warning: Found {n_removed} volumes in diffusion shell (b>50) with zero b-vectors (0,0,0).{RESET}")
+        print(f"{YELLOW}  Excluding invalid volumes to prevent fitting errors...{RESET}")
+        
+        # Create mask of volumes to keep
+        keep_mask = ~bad_indices
+        
+        # Apply filtering
+        dwi_data = dwi_data[..., keep_mask]
+        bvals = bvals[keep_mask]
+        bvecs = bvecs[:, keep_mask]
+        
+        print(f"  New DWI shape after filtering: {dwi_data.shape}")
+        print(f"  New B-values count: {len(bvals)}")
+
     # Count b0 volumes
     b0_count = np.sum(bvals <= 50)
     dwi_count = len(bvals) - b0_count

micaflow/scripts/synth_b0.py

@@ -656,26 +656,6 @@ def main():
         os.makedirs(temp_dir, exist_ok=True)
         print(f"\n{CYAN}Temporary directory:{RESET} {temp_dir}")
 
-        # --- FIX START: Handle 4D input dimensions for SynB0-DISCO compatibility ---
-        # Ensure inputs are strictly 3D. If 4D with N=1 (e.g. 240x240x220x1), 
-        # squeeze them to prevent NumPy "inhomogeneous shape" errors during inference.
-        
-        t1_img_check = nib.load(t1_path)
-        if len(t1_img_check.shape) > 3 and t1_img_check.shape[3] == 1:
-            print(f"{YELLOW}Warning: T1w input is 4D {t1_img_check.shape}. Squeezing to 3D for inference compatibility...{RESET}")
-            t1_3d_path = os.path.join(temp_dir, "t1_3d_input.nii.gz")
-            # Squeeze data and save to temp file
-            nib.save(nib.Nifti1Image(t1_img_check.get_fdata().squeeze(), t1_img_check.affine), t1_3d_path)
-            t1_path = t1_3d_path
-            
-        b0_img_check = nib.load(b0_path)
-        if len(b0_img_check.shape) > 3 and b0_img_check.shape[3] == 1:
-            print(f"{YELLOW}Warning: B0 input is 4D {b0_img_check.shape}. Squeezing to 3D for inference compatibility...{RESET}")
-            b0_3d_path = os.path.join(temp_dir, "b0_3d_input.nii.gz")
-            nib.save(nib.Nifti1Image(b0_img_check.get_fdata().squeeze(), b0_img_check.affine), b0_3d_path)
-            b0_path = b0_3d_path
-        # --- FIX END ---
-
         # Find models
         models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'models')
         print(f"{CYAN}Models directory:{RESET} {models_dir}")