fix visualization bug

app.py

@@ -346,70 +346,8 @@ def create_single_sample_dataloader(patient_idx: int, task: str):
         raise
 
 
-def get_raw_patient_data(patient_idx: int, task: str):
-    """Get raw patient data for visualization"""
-    try:
-        sample_patients, dataset_root = load_sample_data()
-        
-        if task == 'classification':
-            if 'risk' not in sample_patients:
-                raise ValueError("No risk assessment data found")
-            patient_row = sample_patients['risk'].iloc[patient_idx]
-            # Risk dataset uses 'highb' instead of 'dwi'
-            t2w_path = dataset_root / patient_row["t2w"]
-            dwi_path = dataset_root / patient_row["highb"]
-            adc_path = dataset_root / patient_row["adc"]
-        else:
-            if 'UCL' not in sample_patients:
-                raise ValueError("No UCL segmentation data found")
-            patient_row = sample_patients['UCL'].iloc[patient_idx]
-            # UCL dataset uses 'dwi'
-            t2w_path = dataset_root / patient_row["t2w"]
-            dwi_path = dataset_root / patient_row["dwi"]
-            adc_path = dataset_root / patient_row["adc"]
-        
-        # Read raw data for visualization (without transforms)
-        spatial_index = [2, 1, 0]  # DHW format as in dataset
-        
-        def read_nifti(path):
-            vol = nib.load(path)
-            vol = vol.get_fdata().astype(np.float32).transpose(spatial_index)
-            return torch.from_numpy(vol)
-        
-        t2w = read_nifti(str(t2w_path))
-        dwi = read_nifti(str(dwi_path))
-        adc = read_nifti(str(adc_path))
-        
-        # Load prostate mask from existing file
-        prostate_mask = None
-        prostate_mask_path = t2w_path.parent / "prostate_mask.nii.gz"
-        if prostate_mask_path.exists():
-            prostate_mask = read_nifti(str(prostate_mask_path))
-            logger.info(f"Loaded prostate mask from {prostate_mask_path}")
-        else:
-            logger.warning(f"Prostate mask not found at {prostate_mask_path}")
-        
-        # Load ground truth if available
-        ground_truth = None
-        if task == 'classification' and 'pirads' in patient_row:
-            ground_truth = int(patient_row['pirads']) - 2  # Convert to 0-3 range
-        elif task == 'segmentation' and 'lesion' in patient_row:
-            lesion_path = dataset_root / patient_row["lesion"]
-            ground_truth = read_nifti(str(lesion_path))
-            ground_truth = ground_truth > 0  # Binarize
-        
-        return {
-            't2w': t2w,
-            'dwi': dwi,
-            'adc': adc,
-            'ground_truth': ground_truth,
-            'prostate_mask': prostate_mask,
-            'patient_info': patient_row
-        }
-    
-    except Exception as e:
-        logger.error(f"Error loading raw patient data: {e}")
-        raise
+# Removed get_raw_patient_data function as it's no longer needed
+# All data is now processed through the unified preprocessing pipeline
 
 
 def visualize_multimodal_results(preprocessed_data: Dict, prediction: torch.Tensor, task: str) -> plt.Figure:
@@ -1006,7 +944,9 @@ def run_classification_inference(patient_id: str, progress=gr.Progress()):
         logits = []
         model.eval()
         with torch.no_grad():
-            for idx, (img, gt, pid) in enumerate(data_loader):
+            for idx, data in enumerate(data_loader):
+                # Classification only returns (img, gt, pid)
+                img, gt, pid = data
                 img, gt = img.to(args.device), gt.to(args.device)
                 logit = model(img)
                 logits.append(logit)
@@ -1079,7 +1019,14 @@ def run_segmentation_inference(patient_id: str, progress=gr.Progress()):
         # Exact demo inference loop
         model.eval()
         with torch.no_grad():
-            for idx, (img, gt, pid) in enumerate(data_loader):
+            for idx, data in enumerate(data_loader):
+                # Segmentation returns (img, gt, pid, gland) or (img, gt, pid)
+                if len(data) == 4:
+                    img, gt, pid, gland = data
+                else:
+                    img, gt, pid = data
+                    gland = None
+                
                 img, gt = img.to(args.device), gt.to(args.device)
                 if args.sliding_window:
                     pred = sliding_window_inference(
@@ -1178,346 +1125,111 @@ def run_segmentation_inference(patient_id: str, progress=gr.Progress()):
         return fig, empty_3d, f"Error: {str(e)}", gr.update(maximum=63, value=32)
 
 
-# 请用这个新版本完全替换掉你代码中旧的 get_preprocessed_patient_data 函数
-
 def get_preprocessed_patient_data(patient_idx: int, task: str):
     """
-    Get preprocessed patient data that matches exactly what the model uses.
-    This new version ensures the prostate mask undergoes the IDENTICAL transform
-    as the lesion mask (label) by reusing the dataset's transform pipeline.
+    Get preprocessed patient data that matches exactly what the model uses
+    This ensures spatial consistency between displayed images and predictions
     """
     try:
         logger.info(f"Loading preprocessed patient data for patient {patient_idx}, task {task}")
         
-        # Create dataloader to get the preprocessed data and the transform pipeline
-        data_loader, args, dataset_subset = create_single_sample_dataloader(patient_idx, task)
+        # Create dataloader to get the exact preprocessed data that the model uses
+        data_loader, args, dataset = create_single_sample_dataloader(patient_idx, task)
         
-        # Get the original dataset and its transformations
-        # Note: data_loader wraps a Subset, which wraps the original Dataset (UCLSet)
-        original_dataset = dataset_subset.dataset
-        transforms = original_dataset.transform
-
-        # Also get the original raw data to access the raw prostate mask
-        raw_data = get_raw_patient_data(patient_idx, task)
-        
-        # --- 主逻辑：从dataloader中提取模型输入和真实标签 ---
-        preprocessed_data_from_loader = next(iter(data_loader))
-        img, gt, pid = preprocessed_data_from_loader
-
-        preprocessed_data = {
-            'preprocessed_image': img[0],
-            'patient_id': pid[0] if isinstance(pid, (list, tuple)) else pid,
-            'spatial_shape': img.shape[2:],
-            'args': args
-        }
-
-        # 分离出各个模态
-        if img.shape[1] >= 3:
-            preprocessed_data['t2w_preprocessed'] = img[0, 0]
-            preprocessed_data['dwi_preprocessed'] = img[0, 1]
-            preprocessed_data['adc_preprocessed'] = img[0, 2]
-        else: # Fallback
-            preprocessed_data['t2w_preprocessed'] = img[0, 0]
-            preprocessed_data['dwi_preprocessed'] = img[0, 0]
-            preprocessed_data['adc_preprocessed'] = img[0, 0]
-
-        # 处理真实标签 (lesion mask)
-        if gt is not None and torch.is_tensor(gt) and gt.numel() > 0:
-            gt_tensor = gt[0]
-            if gt_tensor.ndim == 4: # [C, D, H, W]
-                gt_tensor = gt_tensor[0] # [D, H, W]
-            preprocessed_data['ground_truth_preprocessed'] = gt_tensor
-        else:
-            preprocessed_data['ground_truth_preprocessed'] = None
-
-        # --- 新的、更简洁的前列腺蒙版处理逻辑 ---
-        preprocessed_data['prostate_mask_preprocessed'] = None
-        if 'prostate_mask' in raw_data and raw_data['prostate_mask'] is not None:
-            logger.info("Applying identical transforms to prostate mask as the lesion mask...")
-            
-            # 准备一个和Dataset输出结构类似的字典，但把prostate_mask放在'label'键上
-            # 'image'键也需要，因为某些空间变换可能需要参考图像的affine信息
-            prostate_mask_input_dict = {
-                'image': raw_data['t2w'], # 使用原始T2W作为参考图像
-                'label': raw_data['prostate_mask']
-            }
-            
-            # 使用从Dataset中获取的完全相同的transforms管道
-            transformed_dict = transforms(prostate_mask_input_dict)
+        # Extract the preprocessed data from the dataloader
+        for idx, data in enumerate(data_loader):
+            # Handle different return formats based on task
+            if task == 'classification':
+                # Classification returns: (img, gt, pid)
+                img, gt, pid = data
+                gland = None
+            else:
+                # Segmentation returns: (img, gt, pid, gland)
+                if len(data) == 4:
+                    img, gt, pid, gland = data
+                    # Remove batch dimension from gland if it exists
+                    if gland is not None and len(gland) > 0:
+                        gland = gland[0]  # Remove batch dimension
+                else:
+                    # Fallback for old format
+                    img, gt, pid = data
+                    gland = None
             
-            # 提取处理后的蒙版
-            processed_mask = transformed_dict['label']
+            # img is the preprocessed multi-modal image [B, C, D, H, W]
+            # gt is the preprocessed ground truth [B, C, D, H, W] or [B, D, H, W]
+            # gland is the preprocessed prostate mask [B, C, D, H, W] or [B, D, H, W] or None
             
-            # 确保二值化
-            processed_mask = (processed_mask > 0.5).float()
+            preprocessed_data = {
+                'preprocessed_image': img[0],  # Remove batch dimension [C, D, H, W]
+                'preprocessed_gt': gt[0] if gt is not None else None,  # Remove batch dimension
+                'patient_id': pid[0] if isinstance(pid, (list, tuple)) else pid,
+                'spatial_shape': img.shape[2:],  # [D, H, W]
+                'args': args
+            }
             
-            non_zero_after = torch.count_nonzero(processed_mask)
-            logger.info(f"Prostate mask after preprocessing: shape={processed_mask.shape}, non-zero voxels={non_zero_after}")
+            # Extract individual modalities from the preprocessed image
+            # Assuming the model input has 3 channels: [T2W, DWI, ADC]
+            if img.shape[1] >= 3:  # Check if we have at least 3 channels
+                preprocessed_data['t2w_preprocessed'] = img[0, 0]  # [D, H, W]
+                preprocessed_data['dwi_preprocessed'] = img[0, 1]  # [D, H, W] 
+                preprocessed_data['adc_preprocessed'] = img[0, 2]  # [D, H, W]
+            else:
+                logger.warning(f"Unexpected number of input channels: {img.shape[1]}")
+                # Fallback: use the same channel for all modalities
+                preprocessed_data['t2w_preprocessed'] = img[0, 0]
+                preprocessed_data['dwi_preprocessed'] = img[0, 0]
+                preprocessed_data['adc_preprocessed'] = img[0, 0]
             
-            if non_zero_after > 0:
-                preprocessed_data['prostate_mask_preprocessed'] = processed_mask
-                logger.info("✅ Successfully applied consistent preprocessing to prostate mask.")
+            # Convert ground truth to proper format
+            if gt is not None:
+                if len(gt.shape) == 4:  # [B, D, H, W]
+                    preprocessed_data['ground_truth_preprocessed'] = gt[0]  # [D, H, W]
+                elif len(gt.shape) == 5:  # [B, C, D, H, W]
+                    preprocessed_data['ground_truth_preprocessed'] = gt[0, 0]  # [D, H, W]
+                else:
+                    logger.warning(f"Unexpected ground truth shape: {gt.shape}")
+                    preprocessed_data['ground_truth_preprocessed'] = None
             else:
-                logger.warning("⚠️ Preprocessing still resulted in an empty prostate mask, even with the correct pipeline.")
-        else:
-            logger.warning("No raw prostate mask found.")
+                preprocessed_data['ground_truth_preprocessed'] = None
             
-        logger.info(f"Successfully loaded and processed data with spatial shape: {preprocessed_data['spatial_shape']}")
+            # Handle prostate gland mask - now from preprocessed data
+            if gland is not None:
+                # Convert gland mask to proper format
+                if len(gland.shape) == 3:  # [D, H, W]
+                    preprocessed_gland = gland
+                elif len(gland.shape) == 4:  # [C, D, H, W]
+                    preprocessed_gland = gland[0]  # [D, H, W]
+                else:
+                    logger.warning(f"Unexpected gland mask shape: {gland.shape}")
+                    preprocessed_gland = None
+                
+                if preprocessed_gland is not None:
+                    # Ensure binary values
+                    preprocessed_gland = (preprocessed_gland > 0.5).float()
+                    non_zero_voxels = torch.count_nonzero(preprocessed_gland)
+                    
+                    logger.info(f"✅ Preprocessed prostate gland mask: shape={preprocessed_gland.shape}, non-zero voxels={non_zero_voxels}")
+                    preprocessed_data['prostate_mask_preprocessed'] = preprocessed_gland if non_zero_voxels > 0 else None
+                else:
+                    logger.warning("Could not process gland mask")
+                    preprocessed_data['prostate_mask_preprocessed'] = None
+            else:
+                logger.info("No prostate gland mask found in preprocessed data")
+                preprocessed_data['prostate_mask_preprocessed'] = None
+                
+            break  # Only process the first (and only) sample
+        
+        logger.info(f"Successfully loaded preprocessed data with spatial shape: {preprocessed_data['spatial_shape']}")
         return preprocessed_data
         
     except Exception as e:
         logger.error(f"Error loading preprocessed patient data: {e}")
-        import traceback
-        logger.error(f"Full traceback: {traceback.format_exc()}")
         raise
 
-# def apply_preprocessing_to_mask(prostate_mask: torch.Tensor, args, transforms) -> torch.Tensor:
-#     """
-#     Apply the same preprocessing transformations to prostate mask as used for image data
-#     This ensures consistent spatial alignment between gland mask and lesion predictions
-#     """
-#     try:
-#         logger.info("Applying consistent preprocessing to prostate mask...")
-#         logger.info(f"Input mask: shape={prostate_mask.shape}, non-zero voxels={torch.count_nonzero(prostate_mask)}")
-        
-#         # Create a dummy data dict similar to what the dataset would create
-#         # We need to mimic the dataset structure for transforms to work
-#         mask_data = {
-#             'image': prostate_mask.unsqueeze(0),  # Add channel dimension [C, D, H, W]
-#             'label': prostate_mask.unsqueeze(0)   # Use same mask as label for processing
-#         }
-        
-#         logger.info(f"Created data dict with shapes: image={mask_data['image'].shape}, label={mask_data['label'].shape}")
-        
-#         # Apply the same transforms that are used for the actual data
-#         try:
-#             logger.info("Applying dataset transforms...")
-#             transformed_data = transforms(mask_data)
-#             preprocessed_mask = transformed_data['label'][0]  # Remove channel dimension
-            
-#             non_zero_after = torch.count_nonzero(preprocessed_mask)
-#             non_zero_original = torch.count_nonzero(prostate_mask)
-#             logger.info(f"Transform result: shape={preprocessed_mask.shape}, non-zero voxels={non_zero_after}")
-#             logger.info(f"Original mask had {non_zero_original} non-zero voxels")
-            
-#             # Check if transforms removed all data
-#             if non_zero_after == 0 and non_zero_original > 0:
-#                 logger.warning("Transforms removed all mask data, falling back to manual preprocessing")
-#                 return apply_manual_preprocessing_to_mask(prostate_mask, args)
-            
-#             # Ensure binary values
-#             preprocessed_mask = (preprocessed_mask > 0.5).float()
-#             final_non_zero = torch.count_nonzero(preprocessed_mask)
-            
-#             logger.info(f"✅ Successfully applied preprocessing to prostate mask: shape={preprocessed_mask.shape}, non-zero voxels={final_non_zero}")
-#             return preprocessed_mask
-            
-#         except Exception as e:
-#             logger.warning(f"Could not apply transforms to prostate mask: {e}")
-#             # Fallback to manual preprocessing
-#             logger.info("Falling back to manual preprocessing...")
-#             return apply_manual_preprocessing_to_mask(prostate_mask, args)
-    
-#     except Exception as e:
-#         logger.error(f"Error applying preprocessing to prostate mask: {e}")
-#         import traceback
-#         logger.error(f"Full traceback: {traceback.format_exc()}")
-        
-#         # Final fallback - just return the original mask if everything fails
-#         logger.warning("Returning original mask as final fallback")
-#         return prostate_mask
-
-# def apply_preprocessing_to_mask(prostate_mask: torch.Tensor, args, transforms) -> torch.Tensor:
-#     """
-#     Apply ONLY SPATIAL preprocessing transformations to the prostate mask.
-#     This ensures consistent spatial alignment without corrupting the mask's binary values.
-#     """
-#     try:
-#         logger.info("Applying consistent SPATIAL preprocessing to prostate mask...")
-#         logger.info(f"Input mask: shape={prostate_mask.shape}, non-zero voxels={torch.count_nonzero(prostate_mask)}")
-
-#         # Define MONAI spatial transforms to filter for
-
-#         # Filter for spatial transforms ONLY
-#         spatial_transforms = [
-#             t for t in transforms.transforms 
-#             if isinstance(t, (Spacing, CenterSpatialCrop, SpatialPad, Resize))
-#         ]
-
-#         if not spatial_transforms:
-#             logger.warning("No spatial transforms found in the pipeline. Falling back to manual preprocessing.")
-#             return apply_manual_preprocessing_to_mask(prostate_mask, args)
 
-#         # Create a new pipeline with only the spatial transforms
-#         # Important: MONAI transforms often expect dictionary inputs
-#         spatial_pipeline = Compose(spatial_transforms)
-        
-#         # Create a data dictionary. Use a key like 'label' or 'mask'
-#         # to avoid intensity normalization which is often applied to 'image' key.
-#         mask_data = {'label': prostate_mask.unsqueeze(0)} # Add channel dim: [1, D, H, W]
-
-#         # Apply the filtered spatial transformations
-#         transformed_data = spatial_pipeline(mask_data)
-#         preprocessed_mask = transformed_data['label'][0] # Remove channel dim
-
-#         # Ensure mask remains binary after interpolation
-#         preprocessed_mask = (preprocessed_mask > 0.5).float()
-
-#         final_non_zero = torch.count_nonzero(preprocessed_mask)
-#         logger.info(f"✅ Successfully applied SPATIAL preprocessing: shape={preprocessed_mask.shape}, non-zero voxels={final_non_zero}")
-
-#         if final_non_zero == 0 and torch.count_nonzero(prostate_mask) > 0:
-#             logger.warning("Spatial transforms resulted in an empty mask. Trying manual fallback.")
-#             return apply_manual_preprocessing_to_mask(prostate_mask, args)
-
-#         return preprocessed_mask
-
-#     except Exception as e:
-#         logger.error(f"Error applying filtered spatial preprocessing: {e}. Falling back to manual method.")
-#         import traceback
-#         logger.error(f"Full traceback: {traceback.format_exc()}")
-#         return apply_manual_preprocessing_to_mask(prostate_mask, args)
-
-# def apply_manual_preprocessing_to_mask(prostate_mask: torch.Tensor, args) -> torch.Tensor:
-#     """
-#     Apply manual preprocessing to prostate mask to match the spatial dimensions
-#     This includes the key transformations: resampling, cropping, and padding
-#     """
-#     try:
-#         logger.info("Applying manual preprocessing to prostate mask...")
-#         logger.info(f"Input mask: shape={prostate_mask.shape}, non-zero voxels={torch.count_nonzero(prostate_mask)}")
-        
-#         # Add batch and channel dimensions for MONAI transforms
-#         mask_tensor = prostate_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, D, H, W]
-#         logger.info(f"After adding dimensions: {mask_tensor.shape}")
-        
-#         # Apply spatial resampling to match the target spacing - be conservative with large masks
-#         if hasattr(args, 'spacing') and args.spacing is not None:
-#             logger.info(f"Applying spacing transform with spacing: {args.spacing}")
-#             # For large original masks, use more conservative resampling
-#             original_size = prostate_mask.shape
-#             if max(original_size) > 300:  # Large original mask
-#                 logger.info("Large original mask detected, using conservative resampling")
-#                 # Use linear interpolation for large masks to preserve more structure
-#                 spacing_transform = Spacing(pixdim=args.spacing, mode='bilinear')
-#             else:
-#                 spacing_transform = Spacing(pixdim=args.spacing, mode='nearest')
-            
-#             mask_tensor_before = mask_tensor.clone()
-#             mask_tensor = spacing_transform(mask_tensor)
-            
-#             logger.info(f"After spacing transform: {mask_tensor.shape}")
-#             logger.info(f"Non-zero voxels after spacing: {torch.count_nonzero(mask_tensor)}")
-            
-#             # Check if spacing transform removed all data
-#             if torch.count_nonzero(mask_tensor) == 0 and torch.count_nonzero(mask_tensor_before) > 0:
-#                 logger.warning("Spacing transform removed all mask data, reverting")
-#                 mask_tensor = mask_tensor_before
-        
-#         # Apply cropping to match the input size - be conservative to preserve mask content
-#         if hasattr(args, 'crop_spatial_size') and args.crop_spatial_size is not None:
-#             current_shape = mask_tensor.shape[2:]  # [D, H, W]
-#             target_shape = args.crop_spatial_size
-            
-#             logger.info(f"Current shape after spacing: {current_shape}, target crop size: {target_shape}")
-            
-#             # Check if we need to crop at all
-#             needs_crop = any(c > t for c, t in zip(current_shape, target_shape))
-            
-#             if needs_crop:
-#                 logger.info(f"Applying center crop to size: {target_shape}")
-                
-#                 # For large reductions, be more conservative
-#                 crop_size_ratios = [t/c for c, t in zip(current_shape, target_shape)]
-#                 min_ratio = min(crop_size_ratios)
-                
-#                 if min_ratio < 0.6:  # Very aggressive crop (>40% reduction)
-#                     logger.warning(f"Very aggressive crop detected (min ratio: {min_ratio:.2f}), using conservative approach")
-#                     # Use a safer crop size that's not too aggressive
-#                     safe_crop_size = tuple(min(c, max(t, int(c * 0.7))) for c, t in zip(current_shape, target_shape))
-#                     logger.info(f"Using safer crop size: {safe_crop_size}")
-#                     crop_transform = CenterSpatialCrop(roi_size=safe_crop_size)
-#                 else:
-#                     crop_transform = CenterSpatialCrop(roi_size=target_shape)
-                
-#                 mask_tensor_before = mask_tensor.clone()
-#                 mask_tensor = crop_transform(mask_tensor)
-                
-#                 logger.info(f"After crop: {mask_tensor.shape}")
-#                 logger.info(f"Non-zero voxels after crop: {torch.count_nonzero(mask_tensor)}")
-                
-#                 # Check if cropping removed all data
-#                 if torch.count_nonzero(mask_tensor) == 0 and torch.count_nonzero(mask_tensor_before) > 0:
-#                     logger.warning("Crop removed all mask data, skipping crop")
-#                     mask_tensor = mask_tensor_before
-#             else:
-#                 logger.info("No cropping needed, current size is within target")
-        
-#         # Apply padding if needed to match exact target size  
-#         if hasattr(args, 'crop_spatial_size') and args.crop_spatial_size is not None:
-#             current_shape = mask_tensor.shape[2:]  # [D, H, W]
-#             target_shape = args.crop_spatial_size
-            
-#             if current_shape != target_shape:
-#                 logger.info(f"Applying padding from {current_shape} to {target_shape}")
-#                 pad_transform = SpatialPad(spatial_size=target_shape, mode='constant')
-#                 mask_tensor = pad_transform(mask_tensor)
-#                 logger.info(f"After padding: {mask_tensor.shape}")
-#                 logger.info(f"Non-zero voxels after padding: {torch.count_nonzero(mask_tensor)}")
-        
-#         # Remove batch and channel dimensions
-#         preprocessed_mask = mask_tensor[0, 0]  # [D, H, W]
-        
-#         # Ensure binary values (thresholding at 0.5)
-#         preprocessed_mask = (preprocessed_mask > 0.5).float()
-        
-#         logger.info(f"Manual preprocessing completed: shape={preprocessed_mask.shape}, non-zero voxels={torch.count_nonzero(preprocessed_mask)}")
-        
-#         return preprocessed_mask
-        
-#     except Exception as e:
-#         logger.error(f"Error in manual preprocessing: {e}")
-#         import traceback
-#         logger.error(f"Full traceback: {traceback.format_exc()}")
-        
-#         # Final fallback - simple resize
-#         logger.warning("Falling back to simple resize")
-#         target_shape = args.crop_spatial_size if hasattr(args, 'crop_spatial_size') else (64, 224, 224)
-        
-#         try:
-#             if prostate_mask.shape != target_shape:
-#                 logger.info(f"Simple resize from {prostate_mask.shape} to {target_shape}")
-#                 zoom_factors = tuple(target_shape[i] / prostate_mask.shape[i] for i in range(3))
-#                 logger.info(f"Zoom factors: {zoom_factors}")
-                
-#                 # Use order=1 (linear) for better preservation of structures, then threshold
-#                 mask_np = prostate_mask.cpu().numpy().astype(np.float32)
-#                 resized_mask = ndimage.zoom(mask_np, zoom_factors, order=1, prefilter=False)
-                
-#                 # Apply threshold to maintain binary nature
-#                 resized_mask = (resized_mask > 0.3).astype(np.float32)  # Lower threshold to preserve more
-                
-#                 result = torch.from_numpy(resized_mask)
-                
-#                 logger.info(f"Simple resize result: shape={result.shape}, non-zero voxels={torch.count_nonzero(result)}")
-                
-#                 if torch.count_nonzero(result) == 0:
-#                     logger.warning("Simple resize resulted in empty mask, trying with very low threshold")
-#                     # Try with an even lower threshold
-#                     resized_mask_low = ndimage.zoom(mask_np, zoom_factors, order=1, prefilter=False)
-#                     resized_mask_low = (resized_mask_low > 0.1).astype(np.float32)
-#                     result = torch.from_numpy(resized_mask_low)
-#                     logger.info(f"Low threshold result: non-zero voxels={torch.count_nonzero(result)}")
-                
-#                 return result
-#             else:
-#                 # Ensure binary
-#                 result = (prostate_mask > 0.5).float()
-#                 return result
-#         except Exception as final_error:
-#             logger.error(f"Simple resize failed: {final_error}")
-#             # Return empty mask if everything fails
-#             return torch.zeros(target_shape, dtype=torch.float32)
+# Removed manual prostate mask preprocessing functions
+# These are no longer needed as the gland mask is now processed 
+# together with image and lesion data in the same pipeline
 
 
 def create_interface():

dataset/dataset_seg.py

@@ -73,14 +73,44 @@ class UCLSet(BaseVolumeDataset):
         img = torch.stack([t2w, dwi, adc], 0)
         seg = self.read(path["lesion"]).unsqueeze(0)
         seg = seg > 0
-        # print(img.shape)
-        # seg = (seg == self.target_class).float()
+        
+        # Load prostate gland mask if available
+        gland = None
+        try:
+            # Try to find prostate_mask.nii.gz in the same directory as t2w
+            t2w_path = os.path.join(self.root, path["t2w"])
+            prostate_mask_path = os.path.join(os.path.dirname(t2w_path), "prostate_mask.nii.gz")
+            if os.path.exists(prostate_mask_path):
+                # Read prostate mask relative to the t2w directory
+                relative_mask_path = os.path.relpath(prostate_mask_path, self.root)
+                gland = self.read(relative_mask_path).unsqueeze(0)
+                gland = gland > 0  # Binarize
+                print(f"Loaded prostate mask: {relative_mask_path}, shape: {gland.shape}, non-zero: {torch.count_nonzero(gland)}")
+            else:
+                print(f"Prostate mask not found at: {prostate_mask_path}")
+        except Exception as e:
+            print(f"Failed to load prostate mask for {path['t2w']}: {e}")
+        
         if self.transforms is not None:
-            trans_dict = self.transforms({"image": img, "label": seg})
+            # Include gland in transforms if available
+            if gland is not None:
+                trans_dict = self.transforms({"image": img, "label": seg, "gland": gland})
+            else:
+                trans_dict = self.transforms({"image": img, "label": seg})
+            
             if type(trans_dict) == list:
                 trans_dict = trans_dict[0]
+            
             img, seg = trans_dict["image"], trans_dict["label"]
-        return img, seg, torch.tensor(idx, dtype=torch.long)
+            
+            # Extract processed gland if it was included
+            if gland is not None and "gland" in trans_dict:
+                gland = trans_dict["gland"]
+            else:
+                gland = None
+        
+        # Return gland as part of the data (we'll modify the app to handle this)
+        return img, seg, torch.tensor(idx, dtype=torch.long), gland
 
 # TODO: need to update; unfinished
 """
@@ -204,22 +234,25 @@ def get_transforms(args):
     train_transforms = [
         NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
         RandRotated(
-            keys=["image", "label"],
+            keys=["image", "label", "gland"],
             prob=0.3,
             range_x=30 / 180 * np.pi,
             keep_size=False,
-            mode=["bilinear", "nearest"],
+            mode=["bilinear", "nearest", "nearest"],
+            allow_missing_keys=True,
         ),
         RandZoomd(
-            keys=["image", "label"],
+            keys=["image", "label", "gland"],
             prob=0.3,
             min_zoom=[1, 0.9, 0.9],
             max_zoom=[1, 1.1, 1.1],
-            mode=["trilinear", "nearest"],
+            mode=["trilinear", "nearest", "nearest"],
+            allow_missing_keys=True,
         ),
         SpatialPadd(
-            keys=["image", "label"],
+            keys=["image", "label", "gland"],
             spatial_size=[round(i * 1.2) for i in args.crop_spatial_size],
+            allow_missing_keys=True,
         ),
         # RandCropByPosNegLabeld(
         #     keys=["image", "label"],
@@ -230,11 +263,12 @@ def get_transforms(args):
         #     num_samples=1,
         # ),
         RandSpatialCropd(
-            keys=["image", "label"],
+            keys=["image", "label", "gland"],
             roi_size=args.crop_spatial_size,
             random_size=False,
+            allow_missing_keys=True,
         ),
-        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
+        RandFlipd(keys=["image", "label", "gland"], prob=0.5, spatial_axis=2, allow_missing_keys=True),
         # BinarizeLabeld(keys=["label"])
         RandScaleIntensityd(keys="image", factors=0.1, prob=0.8),
         RandShiftIntensityd(keys="image", offsets=0.1, prob=0.8),
@@ -247,11 +281,13 @@ def get_transforms(args):
         [
             NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
             CenterSpatialCropd(
-                keys=["image", "label"], roi_size=args.crop_spatial_size
+                keys=["image", "label", "gland"], roi_size=args.crop_spatial_size,
+                allow_missing_keys=True,
             ),
             SpatialPadd(
-                keys=["image", "label"],
+                keys=["image", "label", "gland"],
                 spatial_size=[i for i in args.crop_spatial_size],
+                allow_missing_keys=True,
             ),
             # BinarizeLabeld(keys=["label"])
         ]
@@ -260,11 +296,13 @@ def get_transforms(args):
         [
             NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
             CenterSpatialCropd(
-                keys=["image", "label"], roi_size=args.crop_spatial_size
+                keys=["image", "label", "gland"], roi_size=args.crop_spatial_size,
+                allow_missing_keys=True,
             ),
             SpatialPadd(
-                keys=["image", "label"],
+                keys=["image", "label", "gland"],
                 spatial_size=[i for i in args.crop_spatial_size],
+                allow_missing_keys=True,
             ),
             # BinarizeLabeld(keys=["label"])
         ]