nettracer3d 0.8.8__py3-none-any.whl → 0.8.9__py3-none-any.whl

nettracer3d/segmenter.py

@@ -61,6 +61,74 @@ class InteractiveSegmenter:
         self.previous_z_fore = None
         self.previous_z_back = None
 
+    def get_minimal_chunks_for_coordinates_cpu(self, coordinates_by_z):
+        """
+        Get minimal set of 2D chunks needed to cover the given coordinates
+        Uses same chunking logic as create_2d_chunks() - CPU version
+        """
+        MAX_CHUNK_SIZE = self.twod_chunk_size
+        needed_chunks = {}
+        
+        for z in coordinates_by_z:
+            y_coords = [coord[0] for coord in coordinates_by_z[z]]
+            x_coords = [coord[1] for coord in coordinates_by_z[z]]
+            
+            # Find bounding box of coordinates in this Z-slice
+            y_min, y_max = min(y_coords), max(y_coords)
+            x_min, x_max = min(x_coords), max(x_coords)
+            
+            # Create chunks using same logic as create_2d_chunks
+            y_dim = self.image_3d.shape[1]
+            x_dim = self.image_3d.shape[2]
+            total_pixels = y_dim * x_dim
+            
+            if total_pixels <= MAX_CHUNK_SIZE:
+                # Single chunk for entire Z slice
+                needed_chunks[z] = [[z, 0, y_dim, 0, x_dim]]
+            else:
+                # Multiple chunks - find which ones contain our coordinates
+                largest_dim = 'y' if y_dim >= x_dim else 'x'
+                num_divisions = int(np.ceil(total_pixels / MAX_CHUNK_SIZE))
+                
+                chunks_for_z = []
+                
+                if largest_dim == 'y':
+                    div_size = int(np.ceil(y_dim / num_divisions))
+                    for i in range(0, y_dim, div_size):
+                        end_i = min(i + div_size, y_dim)
+                        # Check if this chunk contains any of our coordinates
+                        if any(i <= y <= end_i-1 for y in y_coords):
+                            chunks_for_z.append([z, i, end_i, 0, x_dim])
+                else:
+                    div_size = int(np.ceil(x_dim / num_divisions))
+                    for i in range(0, x_dim, div_size):
+                        end_i = min(i + div_size, x_dim)
+                        # Check if this chunk contains any of our coordinates
+                        if any(i <= x <= end_i-1 for x in x_coords):
+                            chunks_for_z.append([z, 0, y_dim, i, end_i])
+                
+                needed_chunks[z] = chunks_for_z
+        
+        return needed_chunks
+
+    def compute_features_for_chunk_2d_cpu(self, chunk_coords, speed):
+        """
+        Compute features for a 2D chunk - CPU version
+        chunk_coords: [z, y_start, y_end, x_start, x_end]
+        """
+        z, y_start, y_end, x_start, x_end = chunk_coords
+        
+        # Extract 2D subarray for this chunk
+        subarray_2d = self.image_3d[z, y_start:y_end, x_start:x_end]
+        
+        # Compute features for this chunk
+        if speed:
+            feature_map = self.compute_feature_maps_cpu_2d(image_2d=subarray_2d)
+        else:
+            feature_map = self.compute_deep_feature_maps_cpu_2d(image_2d=subarray_2d)
+        
+        return feature_map, (y_start, x_start)  # Return offset for coordinate mapping
+
     def compute_deep_feature_maps_cpu_2d(self, z=None, image_2d = None):
         """Vectorized detailed version with Gaussian gradient magnitudes, Laplacians, and largest Hessian eigenvalue for 2D images"""
         if z is None:
@@ -131,7 +199,7 @@ class InteractiveSegmenter:
             features[..., feature_idx] = ndimage.laplace(gaussian_img, mode='reflect')
             feature_idx += 1
         
-        # Largest Hessian eigenvalue for each sigma (fully vectorized, 2D version)
+        # Largest Hessian eigenvalue for each sigma (analytical 2D version)
         for sigma in self.sigmas:
             gaussian_img = gaussian_cache[sigma]
             
@@ -140,29 +208,25 @@ class InteractiveSegmenter:
             hyy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[2, 0], mode='reflect')
             hxy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[1, 1], mode='reflect')
             
-            # Vectorized eigenvalue computation using numpy broadcasting
-            # Create arrays with shape (d0, d1, 2, 2) for all 2D Hessian matrices
-            shape = image_2d.shape
-            hessian_matrices = np.zeros(shape + (2, 2))
+            # Analytical eigenvalue computation for 2x2 symmetric matrices
+            # For matrix [[hxx, hxy], [hxy, hyy]], eigenvalues are:
+            # λ = (trace ± sqrt(trace² - 4*det)) / 2
             
-            # Fill the symmetric 2D Hessian matrices
-            hessian_matrices[..., 0, 0] = hxx
-            hessian_matrices[..., 1, 1] = hyy
-            hessian_matrices[..., 0, 1] = hessian_matrices[..., 1, 0] = hxy
+            trace = hxx + hyy
+            det = hxx * hyy - hxy * hxy
             
-            # Reshape for batch eigenvalue computation
-            original_shape = hessian_matrices.shape[:-2]  # (d0, d1)
-            batch_size = np.prod(original_shape)
-            hessian_batch = hessian_matrices.reshape(batch_size, 2, 2)
+            # Calculate discriminant and ensure it's non-negative
+            discriminant = trace * trace - 4 * det
+            discriminant = np.maximum(discriminant, 0)  # Handle numerical errors
             
-            # Compute eigenvalues for all matrices at once
-            eigenvalues_batch = np.real(np.linalg.eigvals(hessian_batch))
+            sqrt_discriminant = np.sqrt(discriminant)
             
-            # Get only the largest eigenvalue for each matrix
-            largest_eigenvalues = np.max(eigenvalues_batch, axis=1)
+            # Calculate both eigenvalues
+            eigenval1 = (trace + sqrt_discriminant) / 2
+            eigenval2 = (trace - sqrt_discriminant) / 2
             
-            # Reshape back to original spatial dimensions
-            largest_eigenvalues = largest_eigenvalues.reshape(original_shape)
+            # Take the larger eigenvalue (most positive/least negative)
+            largest_eigenvalues = np.maximum(eigenval1, eigenval2)
             
             # Add the largest eigenvalue as a feature
             features[..., feature_idx] = largest_eigenvalues
@@ -548,73 +612,98 @@ class InteractiveSegmenter:
         return dict(z_dict)  # Convert back to regular dict
 
     def process_chunk(self, chunk_coords):
-        """
-        Vectorized process_chunk that releases GIL more effectively
-        """
+        """Updated process_chunk with chunked 2D feature computation"""
+        
         if self.realtimechunks is None:
-            # Generate coordinates using vectorized operations
+            # 3D processing (unchanged)
             z_min, z_max = chunk_coords[0], chunk_coords[1]
             y_min, y_max = chunk_coords[2], chunk_coords[3]
             x_min, x_max = chunk_coords[4], chunk_coords[5]
 
-            # More efficient coordinate generation
             z_range = np.arange(z_min, z_max)
             y_range = np.arange(y_min, y_max)
             x_range = np.arange(x_min, x_max)
             
-            # Create coordinate grid efficiently
             z_grid, y_grid, x_grid = np.meshgrid(z_range, y_range, x_range, indexing='ij')
             chunk_coords_array = np.column_stack([
-                z_grid.ravel(), 
-                y_grid.ravel(), 
-                x_grid.ravel()
+                z_grid.ravel(), y_grid.ravel(), x_grid.ravel()
             ])
-        else:
-            # Convert to numpy array for vectorized operations
-            chunk_coords_array = np.array(chunk_coords)
-            z_coords, y_coords, x_coords = chunk_coords_array[:, 0], chunk_coords_array[:, 1], chunk_coords_array[:, 2]
-            z_min, z_max = z_coords.min(), z_coords.max()
-            y_min, y_max = y_coords.min(), y_coords.max()
-            x_min, x_max = x_coords.min(), x_coords.max()
-
-        # Extract subarray
-        subarray = self.image_3d[z_min:z_max+1, y_min:y_max+1, x_min:x_max+1]
+            
+            # Extract subarray
+            subarray = self.image_3d[z_min:z_max+1, y_min:y_max+1, x_min:x_max+1]
+            
+            # Compute features for entire subarray
+            if self.speed:
+                feature_map = self.compute_feature_maps_cpu(subarray)
+            else:
+                feature_map = self.compute_deep_feature_maps_cpu(subarray)
+            
+            # Vectorized feature extraction
+            local_coords = chunk_coords_array - np.array([z_min, y_min, x_min])
+            features = feature_map[local_coords[:, 0], local_coords[:, 1], local_coords[:, 2]]
+            
+            # Vectorized predictions
+            predictions = self.model.predict(features)
+            predictions = np.array(predictions, dtype=bool)
+            
+            # Use boolean indexing to separate coordinates
+            foreground_coords = chunk_coords_array[predictions]
+            background_coords = chunk_coords_array[~predictions]
+            
+            # Convert to sets
+            foreground = set(map(tuple, foreground_coords))
+            background = set(map(tuple, background_coords))
+            
+            return foreground, background
         
-        # Compute features for entire subarray at once
-        if self.speed:
-            feature_map = self.compute_feature_maps_cpu(subarray)
         else:
-            feature_map = self.compute_deep_feature_maps_cpu(subarray)
-        
-        # Vectorized feature extraction
-        # Convert global coordinates to local coordinates in one operation
-        local_coords = chunk_coords_array - np.array([z_min, y_min, x_min])
-        
-        # Extract all features at once using advanced indexing
-        features = feature_map[local_coords[:, 0], local_coords[:, 1], local_coords[:, 2]]
-        
-        # Vectorized predictions (assuming your model can handle batch predictions)
-        if hasattr(self.model, 'predict_batch') or features.ndim > 1:
-            # If model supports batch prediction
+            # 2D processing - compute features for chunk only
+            chunk_coords_array = np.array(chunk_coords)
+            z_coords = chunk_coords_array[:, 0]
+            y_coords = chunk_coords_array[:, 1]
+            x_coords = chunk_coords_array[:, 2]
+            
+            z = int(np.unique(z_coords)[0])  # All coordinates should have same Z
+            
+            # Get chunk bounds
+            y_min, y_max = int(np.min(y_coords)), int(np.max(y_coords))
+            x_min, x_max = int(np.min(x_coords)), int(np.max(x_coords))
+            
+            # Expand bounds slightly to ensure we capture the chunk properly
+            y_min = max(0, y_min)
+            x_min = max(0, x_min)
+            y_max = min(self.image_3d.shape[1], y_max + 1)
+            x_max = min(self.image_3d.shape[2], x_max + 1)
+            
+            # Extract 2D subarray for this chunk
+            subarray_2d = self.image_3d[z, y_min:y_max, x_min:x_max]
+            
+            # Compute features for just this chunk
+            if self.speed:
+                feature_map = self.compute_feature_maps_cpu_2d(image_2d=subarray_2d)
+            else:
+                feature_map = self.compute_deep_feature_maps_cpu_2d(image_2d=subarray_2d)
+            
+            # Convert global coordinates to local chunk coordinates
+            local_y_coords = y_coords - y_min
+            local_x_coords = x_coords - x_min
+            
+            # Extract features using local coordinates
+            features = feature_map[local_y_coords, local_x_coords]
+            
+            # Vectorized predictions
             predictions = self.model.predict(features)
-        else:
-            # Fallback to individual predictions but still vectorized preparation
-            predictions = np.array([self.model.predict([feat]) for feat in features])
-        
-        # Vectorized coordinate assignment
-        predictions = np.array(predictions, dtype=bool)
-        foreground_mask = predictions
-        background_mask = ~predictions
-        
-        # Use boolean indexing to separate coordinates
-        foreground_coords = chunk_coords_array[foreground_mask]
-        background_coords = chunk_coords_array[background_mask]
-        
-        # Convert to sets (still needed for your return format)
-        foreground = set(map(tuple, foreground_coords))
-        background = set(map(tuple, background_coords))
-    
-        return foreground, background
+            predictions = np.array(predictions, dtype=bool)
+            
+            # Use boolean indexing to separate coordinates
+            foreground_coords = chunk_coords_array[predictions]
+            background_coords = chunk_coords_array[~predictions]
+            
+            # Convert to sets
+            foreground = set(map(tuple, foreground_coords))
+            background = set(map(tuple, background_coords))
+            
+            return foreground, background
 
     def twodim_coords(self, y_dim, x_dim, z, chunk_size = None, subrange = None):
 
@@ -967,51 +1056,53 @@ class InteractiveSegmenter:
 
     def get_realtime_chunks_2d(self, chunk_size=None):
         """
-        Create square chunks with 1 z-thickness (2D chunks across XY planes)
+        Updated 2D chunking to match create_2d_chunks logic
         """
         
-        if chunk_size is None:
-            chunk_size = int(np.sqrt(self.twod_chunk_size))
-
-        # Determine if we need to chunk XY planes
-        small_dims = (self.image_3d.shape[1] <= chunk_size and 
-                     self.image_3d.shape[2] <= chunk_size)
-        few_z = self.image_3d.shape[0] <= 100  # arbitrary threshold
-        
-        # If small enough, each Z is one chunk
-        if small_dims and few_z:
-            chunk_size_xy = max(self.image_3d.shape[1], self.image_3d.shape[2])
-        else:
-            chunk_size_xy = chunk_size
-        
-        # Calculate chunks for XY plane
-        y_chunks = (self.image_3d.shape[1] + chunk_size_xy - 1) // chunk_size_xy
-        x_chunks = (self.image_3d.shape[2] + chunk_size_xy - 1) // chunk_size_xy
+        MAX_CHUNK_SIZE = self.twod_chunk_size
         
         # Populate chunk dictionary
         chunk_dict = {}
         
-        # Create chunks for each Z plane (single Z thickness)
+        # Create chunks for each Z plane using the same logic as create_2d_chunks
         for z in range(self.image_3d.shape[0]):
-            if small_dims:
+            y_dim = self.image_3d.shape[1]
+            x_dim = self.image_3d.shape[2]
+            total_pixels = y_dim * x_dim
+            
+            if total_pixels <= MAX_CHUNK_SIZE:
+                # Single chunk for entire Z slice
                 chunk_dict[(z, 0, 0)] = {
-                    'coords': [0, self.image_3d.shape[1], 0, self.image_3d.shape[2]],
+                    'coords': [0, y_dim, 0, x_dim],  # [y_start, y_end, x_start, x_end]
                     'processed': False,
-                    'z': z  # Keep for backward compatibility
+                    'z': z
                 }
             else:
-                # Multiple chunks per Z plane
-                for y_chunk in range(y_chunks):
-                    for x_chunk in range(x_chunks):
-                        y_start = y_chunk * chunk_size_xy
-                        x_start = x_chunk * chunk_size_xy
-                        y_end = min(y_start + chunk_size_xy, self.image_3d.shape[1])
-                        x_end = min(x_start + chunk_size_xy, self.image_3d.shape[2])
-                        
-                        chunk_dict[(z, y_start, x_start)] = {
-                            'coords': [y_start, y_end, x_start, x_end],
+                # Multiple chunks per Z plane - divide along largest dimension
+                largest_dim = 'y' if y_dim >= x_dim else 'x'
+                num_divisions = int(np.ceil(total_pixels / MAX_CHUNK_SIZE))
+                
+                if largest_dim == 'y':
+                    # Divide along Y dimension
+                    div_size = int(np.ceil(y_dim / num_divisions))
+                    for i in range(0, y_dim, div_size):
+                        end_i = min(i + div_size, y_dim)
+                        # Use (z, y_start, x_start) as key for consistency
+                        chunk_dict[(z, i, 0)] = {
+                            'coords': [i, end_i, 0, x_dim],  # [y_start, y_end, x_start, x_end]
                             'processed': False,
-                            'z': z  # Keep for backward compatibility
+                            'z': z
+                        }
+                else:
+                    # Divide along X dimension
+                    div_size = int(np.ceil(x_dim / num_divisions))
+                    for i in range(0, x_dim, div_size):
+                        end_i = min(i + div_size, x_dim)
+                        # Use (z, y_start, x_start) as key for consistency
+                        chunk_dict[(z, 0, i)] = {
+                            'coords': [0, y_dim, i, end_i],  # [y_start, y_end, x_start, x_end]
+                            'processed': False,
+                            'z': z
                         }
         
         self.realtimechunks = chunk_dict
@@ -1041,33 +1132,37 @@ class InteractiveSegmenter:
         
         return slice_foreground_features, slice_background_features
 
-    def extract_features_parallel(self, slices: List[int], speed: Any, use_gpu: bool,
-                                 z_fores: Dict[int, List[Tuple[int, int]]], 
-                                 z_backs: Dict[int, List[Tuple[int, int]]]) -> Tuple[List[Any], List[Any]]:
+    def extract_features_parallel(self, needed_chunks, speed, use_gpu, z_fores, z_backs):
         """
-        Process feature extraction using ThreadPoolExecutor for parallel execution.
+        Updated version that processes chunks instead of full Z-slices
         """
         max_cores = multiprocessing.cpu_count()
         foreground_features = []
         background_features = []
         
+        # Flatten all chunks into a single list for parallel processing
+        all_chunk_tasks = []
+        for z in needed_chunks:
+            for chunk_coords in needed_chunks[z]:
+                all_chunk_tasks.append((chunk_coords, z_fores, z_backs, speed))
+        
         with ThreadPoolExecutor(max_workers=max_cores) as executor:
-            # Submit all slice processing tasks
-            future_to_slice = {
-                executor.submit(self.process_slice_features, z, speed, use_gpu, z_fores, z_backs): z 
-                for z in slices
+            # Submit all chunk processing tasks
+            future_to_chunk = {
+                executor.submit(self.process_chunk_features_for_training, chunk_task): chunk_task 
+                for chunk_task in all_chunk_tasks
             }
             
             # Collect results as they complete
-            for future in future_to_slice:
-                slice_foreground, slice_background = future.result()
-                foreground_features.extend(slice_foreground)
-                background_features.extend(slice_background)
+            for future in future_to_chunk:
+                chunk_foreground, chunk_background = future.result()
+                foreground_features.extend(chunk_foreground)
+                background_features.extend(chunk_background)
         
         return foreground_features, background_features
 
-    def segment_volume_realtime(self, gpu = False):
-
+    def segment_volume_realtime(self, gpu=False):
+        """Updated realtime segmentation with chunked 2D processing"""
 
         if self.realtimechunks is None:
             if not self.use_two:
@@ -1075,11 +1170,10 @@ class InteractiveSegmenter:
             else:
                 self.get_realtime_chunks_2d()
         else:
-            for chunk_pos in self.realtimechunks:  # chunk_pos is the (z, y_start, x_start) tuple
+            for chunk_pos in self.realtimechunks:
                 self.realtimechunks[chunk_pos]['processed'] = False
 
         chunk_dict = self.realtimechunks
-
         
         def get_nearest_unprocessed_chunk(self):
             """Get nearest unprocessed chunk prioritizing current Z"""
@@ -1092,7 +1186,6 @@ class InteractiveSegmenter:
                               if pos[0] == curr_z and not info['processed']]
             
             if current_z_chunks:
-                # Find nearest chunk in current Z plane using the chunk positions from the key
                 nearest = min(current_z_chunks, 
                             key=lambda x: ((x[0][1] - curr_y) ** 2 + 
                                          (x[0][2] - curr_x) ** 2))
@@ -1107,7 +1200,6 @@ class InteractiveSegmenter:
             
             if available_z:
                 target_z = available_z[0][0]
-                # Find nearest chunk in target Z plane
                 z_chunks = [(pos, info) for pos, info in chunk_dict.items() 
                            if pos[0] == target_z and not info['processed']]
                 nearest = min(z_chunks, 
@@ -1117,41 +1209,31 @@ class InteractiveSegmenter:
             
             return None
         
-
         while True:
-            # Find nearest unprocessed chunk using class attributes
             chunk_idx = get_nearest_unprocessed_chunk(self)
             if chunk_idx is None:
                 break
                 
-            # Process the chunk directly
             chunk = chunk_dict[chunk_idx]
             chunk['processed'] = True
-            coords = chunk['coords']
+            coords = chunk['coords']  # [y_start, y_end, x_start, x_end]
+            z = chunk['z']
 
-            coords = np.stack(np.meshgrid(
-                [chunk['z']],
+            # Generate coordinates for this chunk
+            coords_array = np.stack(np.meshgrid(
+                [z],
                 np.arange(coords[0], coords[1]),
                 np.arange(coords[2], coords[3]),
                 indexing='ij'
             )).reshape(3, -1).T
 
-            coords = list(map(tuple, coords))
-
+            coords_list = list(map(tuple, coords_array))
             
-            # Process the chunk directly based on whether GPU is available
-            if gpu:
-                try:
-                    fore, back = self.process_chunk_GPU(coords)
-                except:
-                    fore, back = self.process_chunk(coords)
-            else:
-                fore, back = self.process_chunk(coords)
+            # Process the chunk with updated method
+            fore, back = self.process_chunk(coords_list)
             
-            # Yield the results
             yield fore, back
 
-
     def cleanup(self):
         """Clean up GPU memory"""
         if self.use_gpu:
@@ -1245,8 +1327,8 @@ class InteractiveSegmenter:
         
         return foreground_features, background_features
 
-    def train_batch(self, foreground_array, speed = True, use_gpu = False, use_two = False, mem_lock = False, saving = False):
-        """Train directly on foreground and background arrays"""
+    def train_batch(self, foreground_array, speed=True, use_gpu=False, use_two=False, mem_lock=False, saving=False):
+        """Updated train_batch with chunked 2D processing"""
 
         if not saving:
             print("Training model...")
@@ -1268,85 +1350,73 @@ class InteractiveSegmenter:
         self.mem_lock = mem_lock
 
         if use_two:
-
-            #changed = [] #Track which slices need feature maps
-
-            if not self.use_two: #Clarifies if we need to redo feature cache for 2D
+            if not self.use_two:
                 self.use_two = True
-
             self.two_slices = []
 
-
             # Get foreground coordinates and features
             z_fore, y_fore, x_fore = np.where(foreground_array == 1)
-
-
             fore_coords = list(zip(z_fore, y_fore, x_fore))
             
             # Get background coordinates and features
             z_back, y_back, x_back = np.where(foreground_array == 2)
-
             back_coords = list(zip(z_back, y_back, x_back))
-
+            
             foreground_features = []
             background_features = []
-
+            
+            # Organize coordinates by Z
             z_fores = self.organize_by_z(fore_coords)
             z_backs = self.organize_by_z(back_coords)
-            slices = set(list(z_fores.keys()) + list(z_backs.keys()))
-
+            
+            # Combine all Z-slices that have coordinates
+            all_z_coords = {}
+            for z in z_fores:
+                if z not in all_z_coords:
+                    all_z_coords[z] = []
+                all_z_coords[z].extend(z_fores[z])
+            for z in z_backs:
+                if z not in all_z_coords:
+                    all_z_coords[z] = []
+                all_z_coords[z].extend(z_backs[z])
+            
+            # Get minimal chunks needed to cover all coordinates
+            needed_chunks = self.get_minimal_chunks_for_coordinates_cpu(all_z_coords)
+            
+            # Use existing parallel infrastructure with chunked approach
             foreground_features, background_features = self.extract_features_parallel(
-                slices, speed, use_gpu, z_fores, z_backs
+                needed_chunks, speed, use_gpu, z_fores, z_backs
             )
 
-
-        else: #Forces ram efficiency
-
+        else:
+            # 3D processing (unchanged - your existing code)
             box_size = self.master_chunk
-
-            # Memory-efficient approach: compute features only for necessary subarrays
             foreground_features = []
             background_features = []
             
-            # Find coordinates of foreground and background scribbles
             z_fore = np.argwhere(foreground_array == 1)
             z_back = np.argwhere(foreground_array == 2)
             
-            # If no scribbles, return empty lists
             if len(z_fore) == 0 and len(z_back) == 0:
                 return foreground_features, background_features
             
-            # Get dimensions of the input array
             depth, height, width = foreground_array.shape
-            
-            # Determine the minimum number of boxes needed to cover all scribbles
-            half_box = box_size // 2
-            
-            # Step 1: Find the minimum set of boxes that cover all scribbles
-            # We'll divide the volume into a grid of boxes of size box_size
-            
-            # Calculate how many boxes are needed in each dimension
             z_grid_size = (depth + box_size - 1) // box_size
             y_grid_size = (height + box_size - 1) // box_size
             x_grid_size = (width + box_size - 1) // box_size
             
-            # Track which grid cells contain scribbles
             grid_cells_with_scribbles = set()
             
-            # Map original coordinates to grid cells
             for z, y, x in np.vstack((z_fore, z_back)) if len(z_back) > 0 else z_fore:
                 grid_z = z // box_size
                 grid_y = y // box_size
                 grid_x = x // box_size
                 grid_cells_with_scribbles.add((grid_z, grid_y, grid_x))
             
-            # Create a mapping from original coordinates to their corresponding subarray and local coordinates
-            coord_mapping = {}
-            
-            # Step 2: Process each grid cell that contains scribbles
-
-            foreground_features, background_features = self.process_grid_cells_parallel(grid_cells_with_scribbles, box_size, depth, height, width, foreground_array)
+            foreground_features, background_features = self.process_grid_cells_parallel(
+                grid_cells_with_scribbles, box_size, depth, height, width, foreground_array)
 
+        # Rest of the method unchanged (combining with previous features, training, etc.)
         if self.previous_foreground is not None:
             failed = True
             try:
@@ -1371,7 +1441,6 @@ class InteractiveSegmenter:
                 print("Could not combine new model with old loaded model. Perhaps you are trying to combine a quick model with a deep model? I cannot combine these...")
 
         if saving:
-
             return foreground_features, background_features, z_fore, z_back
 
         # Combine features and labels
@@ -1386,13 +1455,8 @@ class InteractiveSegmenter:
             print(y)
 
         self.current_speed = speed
-                
-
-
-
         print("Done")
 
-
     def save_model(self, file_name, foreground_array):
 
         print("Saving model data")
@@ -1453,3 +1517,47 @@ class InteractiveSegmenter:
 
         return output
 
+
+    def process_chunk_features_for_training(self, chunk_task):
+        """
+        Process a single chunk for training feature extraction
+        chunk_task: (chunk_coords, z_fores, z_backs, speed)
+        """
+        chunk_coords, z_fores, z_backs, speed = chunk_task
+        z = chunk_coords[0]
+        
+        # Compute features for this chunk
+        feature_map, (y_offset, x_offset) = self.compute_features_for_chunk_2d_cpu(chunk_coords, speed)
+        
+        chunk_foreground_features = []
+        chunk_background_features = []
+        
+        # Extract foreground features from this chunk
+        if z in z_fores:
+            for y, x in z_fores[z]:
+                # Check if this coordinate is in the current chunk
+                y_start, y_end = chunk_coords[1], chunk_coords[2]
+                x_start, x_end = chunk_coords[3], chunk_coords[4]
+                
+                if y_start <= y < y_end and x_start <= x < x_end:
+                    # Convert global coordinates to local chunk coordinates
+                    local_y = y - y_offset
+                    local_x = x - x_offset
+                    feature_vector = feature_map[local_y, local_x]
+                    chunk_foreground_features.append(feature_vector)
+        
+        # Extract background features from this chunk
+        if z in z_backs:
+            for y, x in z_backs[z]:
+                # Check if this coordinate is in the current chunk
+                y_start, y_end = chunk_coords[1], chunk_coords[2]
+                x_start, x_end = chunk_coords[3], chunk_coords[4]
+                
+                if y_start <= y < y_end and x_start <= x < x_end:
+                    # Convert global coordinates to local chunk coordinates
+                    local_y = y - y_offset
+                    local_x = x - x_offset
+                    feature_vector = feature_map[local_y, local_x]
+                    chunk_background_features.append(feature_vector)
+        
+        return chunk_foreground_features, chunk_background_features