nettracer3d 0.8.4__py3-none-any.whl → 0.8.5__py3-none-any.whl

nettracer3d/segmenter.py

@@ -6,6 +6,7 @@ import threading
 from scipy import ndimage
 import multiprocessing
 from collections import defaultdict
+from typing import List, Dict, Tuple, Any
 
 class InteractiveSegmenter:
     def __init__(self, image_3d, use_gpu=False):
@@ -37,7 +38,6 @@ class InteractiveSegmenter:
         self.two_slices = []
         self.speed = True
         self.cur_gpu = False
-        self.map_slice = None
         self.prev_z = None
         self.previewing = False
 
@@ -48,10 +48,12 @@ class InteractiveSegmenter:
         self.mem_lock = False
 
         #Adjustable feature map params:
-        self.alphas = [1,2,4,8]
+        self.sigmas = [1,2,4,8]
         self.windows = 10
         self.dogs = [(1, 2), (2, 4), (4, 8)]
         self.master_chunk = 49
+        self.twod_chunk_size = 262144
+        self.batch_amplifier = 1
 
         #Data when loading prev model:
         self.previous_foreground = None
@@ -59,618 +61,473 @@ class InteractiveSegmenter:
         self.previous_z_fore = None
         self.previous_z_back = None
 
-    def segment_slice_chunked(self, slice_z, block_size = 49):
-        """
-        A completely standalone method to segment a single z-slice in chunks
-        with improved safeguards.
-        """
-        # Check if we're already processing this slice
-        if self._currently_processing and self._currently_processing == slice_z:
-            return
+    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:
+            z = self.image_3d.shape[0] // 2  # Use middle slice if not specified
         
-        # Set processing flag with the slice we're processing
-        self._currently_processing = slice_z
+        if image_2d is None:
+            image_2d = self.image_3d[z, :, :]
+
+        if image_2d.ndim == 3 and image_2d.shape[-1] == 3:
+            # RGB case - process each channel
+            features_per_channel = []
+            for channel in range(3):
+                channel_features = self.compute_deep_feature_maps_cpu_2d(image_2d = image_2d[..., channel])
+                features_per_channel.append(channel_features)
+            
+            # Stack all channel features
+            return np.concatenate(features_per_channel, axis=-1)
         
-        try:
-            
-            # First attempt to get the feature map
-            feature_map = None
-            
-            try:
-                if slice_z in self.feature_cache:
-                    feature_map = self.feature_cache[slice_z]
-                elif hasattr(self, 'map_slice') and self.map_slice is not None and slice_z == self.current_z:
-                    feature_map = self.map_slice
-                else:
-                    # Generate new feature map
-                    try:
-                        feature_map = self.get_feature_map_slice(slice_z, self.current_speed, False)
-                        self.map_slice = feature_map
-                        # Cache the feature map for future use
-                        #if not hasattr(self, 'feature_cache'):
-                            #self.feature_cache = {}
-                        #self.feature_cache[slice_z] = feature_map
-                    except Exception as e:
-                        print(f"Error generating feature map: {e}")
-                        import traceback
-                        traceback.print_exc()
-                        return  # Exit if we can't generate the feature map
-            except:
-                # Generate new feature map
-                #self.feature_cache = {}
-                try:
-                    feature_map = self.get_feature_map_slice(slice_z, self.current_speed, False)
-                    self.map_slice = feature_map
-                    # Cache the feature map for future use
-                    #if not hasattr(self, 'feature_cache'):
-                        #self.feature_cache = {}
-                    #self.feature_cache[slice_z] = feature_map
-                except Exception as e:
-                    print(f"Error generating feature map: {e}")
-                    import traceback
-                    traceback.print_exc()
-                    return  # Exit if we can't generate the feature map
-
-            
-            # Check that we have a valid feature map
-            if feature_map is None:
-                return
-            
-            # Get dimensions of the slice
-            y_size, x_size = self.image_3d.shape[1], self.image_3d.shape[2]
-            chunk_count = 0
-            
-            # Process in blocks for chunked feedback
-            for y_start in range(0, y_size, block_size):
-                if self._currently_processing != slice_z:
-                    return
-                    
-                for x_start in range(0, x_size, block_size):
-                    if self._currently_processing != slice_z:
-                        return
-                        
-                    y_end = min(y_start + block_size, y_size)
-                    x_end = min(x_start + block_size, x_size)
-                    
-                    # Create coordinates and features for this block
-                    coords = []
-                    features = []
-                    
-                    for y in range(y_start, y_end):
-                        for x in range(x_start, x_end):
-                            coords.append((slice_z, y, x))
-                            features.append(feature_map[y, x])
-                    
-                    # Skip empty blocks
-                    if not coords:
-                        continue
-                    
-                    # Predict
-                    try:
-                        try:
-                            predictions = self.model.predict(features)
-                        except ValueError:
-                            self.feature_cache = None
-                            self.map_slice = None
-                            return None, None
-                        
-                        # Split results
-                        foreground = set()
-                        background = set()
-                        
-                        for coord, pred in zip(coords, predictions):
-                            if pred:
-                                foreground.add(coord)
-                            else:
-                                background.add(coord)
-                        
-                        # Yield this chunk
-                        chunk_count += 1
-                        yield foreground, background
-                        
-                    except Exception as e:
-                        print(f"Error processing chunk: {e}")
-                        import traceback
-                        traceback.print_exc()
-            
         
-        finally:
-            # Only clear if we're still processing the same slice
-            # (otherwise, another slice might have taken over)
-            if self._currently_processing == slice_z:
-                self._currently_processing = None
+        # Calculate total number of features
+        num_basic_features = 1 + len(self.sigmas) + len(self.dogs)  # original + gaussians + dogs
+        num_gradient_features = len(self.sigmas)  # gradient magnitude for each sigma
+        num_laplacian_features = len(self.sigmas)  # laplacian for each sigma
+        num_hessian_features = len(self.sigmas) * 1  # 1 eigenvalue (largest) for each sigma
+        
+        total_features = num_basic_features + num_gradient_features + num_laplacian_features + num_hessian_features
+        
+        # Pre-allocate result array
+        features = np.empty(image_2d.shape + (total_features,), dtype=image_2d.dtype)
+        features[..., 0] = image_2d
+        
+        feature_idx = 1
+        
+        # Cache for Gaussian filters - only compute each sigma once
+        gaussian_cache = {}
+        
+        # Compute all unique sigmas needed (from both sigmas and dogs)
+        all_sigmas = set(self.sigmas)
+        for s1, s2 in self.dogs:
+            all_sigmas.add(s1)
+            all_sigmas.add(s2)
+        
+        # Pre-compute all Gaussian filters
+        for sigma in all_sigmas:
+            gaussian_cache[sigma] = ndimage.gaussian_filter(image_2d, sigma)
+        
+        # Gaussian smoothing - use cached results
+        for sigma in self.sigmas:
+            features[..., feature_idx] = gaussian_cache[sigma]
+            feature_idx += 1
+        
+        # Difference of Gaussians - use cached results
+        for s1, s2 in self.dogs:
+            features[..., feature_idx] = gaussian_cache[s1] - gaussian_cache[s2]
+            feature_idx += 1
+        
+        # Gaussian gradient magnitudes for each sigma (vectorized, 2D version)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            gx = ndimage.sobel(gaussian_img, axis=1, mode='reflect')  # x direction
+            gy = ndimage.sobel(gaussian_img, axis=0, mode='reflect')  # y direction
+            features[..., feature_idx] = np.sqrt(gx**2 + gy**2)
+            feature_idx += 1
+        
+        # Laplacian of Gaussian for each sigma (vectorized, 2D version)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            features[..., feature_idx] = ndimage.laplace(gaussian_img, mode='reflect')
+            feature_idx += 1
+        
+        # Largest Hessian eigenvalue for each sigma (fully vectorized, 2D version)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            
+            # Compute second derivatives (Hessian components) - all vectorized for 2D
+            hxx = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 2], mode='reflect')
+            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))
+            
+            # 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
+            
+            # 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)
+            
+            # Compute eigenvalues for all matrices at once
+            eigenvalues_batch = np.real(np.linalg.eigvals(hessian_batch))
+            
+            # Get only the largest eigenvalue for each matrix
+            largest_eigenvalues = np.max(eigenvalues_batch, axis=1)
+            
+            # Reshape back to original spatial dimensions
+            largest_eigenvalues = largest_eigenvalues.reshape(original_shape)
+            
+            # Add the largest eigenvalue as a feature
+            features[..., feature_idx] = largest_eigenvalues
+            feature_idx += 1
+        
+        # Normalize only morphological features, keep intensity features raw
+        intensity_features = features[..., :num_basic_features]  # original + gaussians + DoGs
+        morphology_features = features[..., num_basic_features:]  # gradients + laplacians + eigenvalues
+
+        # Normalize only morphological features
+        morph_means = np.mean(morphology_features, axis=(0, 1), keepdims=True)
+        morph_stds = np.std(morphology_features, axis=(0, 1), keepdims=True)
+        morph_stds = np.where(morph_stds == 0, 1, morph_stds)
+        morphology_features = (morphology_features - morph_means) / morph_stds
+
+        # Recombine
+        features = np.concatenate([intensity_features, morphology_features], axis=-1)
+        
+        return features
 
 
+    def compute_feature_maps_cpu_2d(self, z=None, image_2d = None):
+        """Compute feature maps for 2D images using CPU with caching optimization"""
+        if image_2d is None:
+            image_2d = self.image_3d[z, :, :]
 
-    def compute_deep_feature_maps_cpu_parallel(self, image_3d=None):
-        """Compute deep feature maps using CPU with thread-based parallelism"""
-        if image_3d is None:
-            image_3d = self.image_3d
+        if image_2d.ndim == 3 and image_2d.shape[-1] == 3:
+            # RGB case - process each channel
+            features_per_channel = []
+            for channel in range(3):
+                channel_features = self.compute_feature_maps_cpu_2d(image_2d = image_2d[..., channel])
+                features_per_channel.append(channel_features)
+            
+            # Stack all channel features
+            return np.concatenate(features_per_channel, axis=-1)
         
-        original_shape = image_3d.shape
+        # Pre-allocate result array
+        num_features = len(self.sigmas) + len(self.dogs) + 2  # +2 for original image + gradient
+        features = np.empty(image_2d.shape + (num_features,), dtype=image_2d.dtype)
         
-        # Use ThreadPoolExecutor for parallelization
-        with ThreadPoolExecutor(max_workers=min(7, multiprocessing.cpu_count())) as executor:
-            # Stage 1: Independent computations that can be parallelized
-            futures = []
-            
-            # Gaussian smoothing
-            def compute_gaussian(sigma):
-                return ndimage.gaussian_filter(image_3d, sigma)
-            
-            for sigma in self.alphas:
-                future = executor.submit(compute_gaussian, sigma)
-                futures.append(('gaussian', sigma, future))
+        # Include original image as first feature
+        features[..., 0] = image_2d
+        feature_idx = 1
+        
+        # Cache for Gaussian filters - only compute each sigma once
+        gaussian_cache = {}
+        
+        # Compute all unique sigmas needed (from both sigmas and dogs)
+        all_sigmas = set(self.sigmas)
+        for s1, s2 in self.dogs:
+            all_sigmas.add(s1)
+            all_sigmas.add(s2)
+        
+        # Pre-compute all Gaussian filters
+        for sigma in all_sigmas:
+            gaussian_cache[sigma] = ndimage.gaussian_filter(image_2d, sigma)
+        
+        # Gaussian smoothing - use cached results
+        for sigma in self.sigmas:
+            features[..., feature_idx] = gaussian_cache[sigma]
+            feature_idx += 1
+        
+        # Difference of Gaussians - use cached results
+        for s1, s2 in self.dogs:
+            features[..., feature_idx] = gaussian_cache[s1] - gaussian_cache[s2]
+            feature_idx += 1
+        
+        # Gradient magnitude (2D version)
+        gx = ndimage.sobel(image_2d, axis=1, mode='reflect')  # x direction
+        gy = ndimage.sobel(image_2d, axis=0, mode='reflect')  # y direction
+        features[..., feature_idx] = np.sqrt(gx**2 + gy**2)
+        
+        return features
 
-            def compute_dog_local(img, s1, s2):
-                g1 = ndimage.gaussian_filter(img, s1)
-                g2 = ndimage.gaussian_filter(img, s2)
-                return g1 - g2
+    def compute_deep_feature_maps_cpu(self, image_3d=None):
+        """Vectorized detailed version with Gaussian gradient magnitudes, Laplacians, and largest Hessian eigenvalue only"""
+        if image_3d is None:
+            image_3d = self.image_3d
 
-            # Difference of Gaussians
-            for (s1, s2) in self.dogs:
-                
-                future = executor.submit(compute_dog_local, image_3d, s1, s2)
-                futures.append(('dog', s1, future))
-            
-            # Local statistics computation
-            def compute_local_mean():
-                window_size = self.windows
-                kernel = np.ones((window_size, window_size, window_size)) / (window_size**3)
-                return ndimage.convolve(image_3d, kernel, mode='reflect')
-            
-            future = executor.submit(compute_local_mean)
-            futures.append(('local_mean', None, future))
-            
-            def compute_local_variance():
-                window_size = self.windows
-                kernel = np.ones((window_size, window_size, window_size)) / (window_size**3)
-                mean = np.mean(image_3d)
-                return ndimage.convolve((image_3d - mean)**2, kernel, mode='reflect')
-            
-            future = executor.submit(compute_local_variance)
-            futures.append(('local_var', None, future))
-            
-            # Gradient computation
-            def compute_gradients():
-                gx = ndimage.sobel(image_3d, axis=2, mode='reflect')
-                gy = ndimage.sobel(image_3d, axis=1, mode='reflect')
-                gz = ndimage.sobel(image_3d, axis=0, mode='reflect')
-                return gx, gy, gz
-            
-            future = executor.submit(compute_gradients)
-            futures.append(('gradients', None, future))
-            
-            # Collect results for the independent computations
-            results = {}
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    if task_type == 'gradients':
-                        # Store the gradient components separately
-                        gx, gy, gz = result
-                        results['gx'] = gx
-                        results['gy'] = gy
-                        results['gz'] = gz
-                    else:
-                        results[f"{task_type}_{params}" if params is not None else task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
-            
-            # Stage 2: Dependent computations that need results from Stage 1
-            futures = []
-            
-            # Gradient magnitude (depends on gradients)
-            def compute_gradient_magnitude(gx, gy, gz):
-                return np.sqrt(gx**2 + gy**2 + gz**2)
-            
-            future = executor.submit(compute_gradient_magnitude, 
-                                   results['gx'], results['gy'], results['gz'])
-            futures.append(('gradient_magnitude', None, future))
-            
-            # Second-order gradients (depend on first gradients)
-            def compute_second_derivatives(gx, gy, gz):
-                gxx = ndimage.sobel(gx, axis=2, mode='reflect')
-                gyy = ndimage.sobel(gy, axis=1, mode='reflect')
-                gzz = ndimage.sobel(gz, axis=0, mode='reflect')
-                return gxx, gyy, gzz
-            
-            future = executor.submit(compute_second_derivatives, 
-                                   results['gx'], results['gy'], results['gz'])
-            futures.append(('second_derivatives', None, future))
-            
-            # Collect results for the dependent computations
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    if task_type == 'second_derivatives':
-                        # Store the second derivative components separately
-                        gxx, gyy, gzz = result
-                        results['gxx'] = gxx
-                        results['gyy'] = gyy
-                        results['gzz'] = gzz
-                    else:
-                        results[task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
-            
-            # Stage 3: Final computations that depend on Stage 2 results
-            futures = []
-            
-            # Laplacian and Hessian determinant (depend on second derivatives)
-            def compute_laplacian(gxx, gyy, gzz):
-                return gxx + gyy + gzz
+        if image_3d.ndim == 4 and image_3d.shape[-1] == 3:
+            # RGB case - process each channel
+            features_per_channel = []
+            for channel in range(3):
+                channel_features = self.compute_deep_feature_maps_cpu(image_3d[..., channel])
+                features_per_channel.append(channel_features)
             
-            future = executor.submit(compute_laplacian, 
-                                   results['gxx'], results['gyy'], results['gzz'])
-            futures.append(('laplacian', None, future))
-            
-            def compute_hessian_det(gxx, gyy, gzz):
-                return gxx * gyy * gzz
-            
-            future = executor.submit(compute_hessian_det, 
-                                   results['gxx'], results['gyy'], results['gzz'])
-            futures.append(('hessian_det', None, future))
-            
-            # Collect final results
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    results[task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
+            # Stack all channel features
+            return np.concatenate(features_per_channel, axis=-1)
         
-        # Organize results in the expected order
-        features = []
         
-        # Add Gaussian features
-        for sigma in self.alphas:
-            features.append(results[f'gaussian_{sigma}'])
-
-        for sigma in self.dogs:
-            features.append(results[f'dog_{sigma[0]}'])
+        # Calculate total number of features
+        num_basic_features = 1 + len(self.sigmas) + len(self.dogs)  # original + gaussians + dogs
+        num_gradient_features = len(self.sigmas)  # gradient magnitude for each sigma
+        num_laplacian_features = len(self.sigmas)  # laplacian for each sigma
+        num_hessian_features = len(self.sigmas) * 1  # 1 eigenvalue (largest) for each sigma
         
-        # Add local statistics
-        features.append(results['local_mean'])
-        features.append(results['local_var'])
+        total_features = num_basic_features + num_gradient_features + num_laplacian_features + num_hessian_features
         
-        # Add gradient magnitude
-        features.append(results['gradient_magnitude'])
+        # Pre-allocate result array
+        features = np.empty(image_3d.shape + (total_features,), dtype=image_3d.dtype)
+        features[..., 0] = image_3d
         
-        # Add Laplacian and Hessian determinant
-        features.append(results['laplacian'])
-        features.append(results['hessian_det'])
+        feature_idx = 1
         
-        # Verify shapes
-        for i, feat in enumerate(features):
-            if feat.shape != original_shape:
-                feat_adjusted = np.expand_dims(feat, axis=0)
-                if feat_adjusted.shape != original_shape:
-                    raise ValueError(f"Feature {i} has shape {feat.shape}, expected {original_shape}")
-                features[i] = feat_adjusted
+        # Cache for Gaussian filters - only compute each sigma once
+        gaussian_cache = {}
         
-        return np.stack(features, axis=-1)
-
-
-    def compute_deep_feature_maps_cpu_2d_parallel(self, z=None):
-        """Compute 2D feature maps using CPU with thread-based parallelism"""
-        image_2d = self.image_3d[z, :, :]
-        original_shape = image_2d.shape
+        # Compute all unique sigmas needed (from both sigmas and dogs)
+        all_sigmas = set(self.sigmas)
+        for s1, s2 in self.dogs:
+            all_sigmas.add(s1)
+            all_sigmas.add(s2)
         
-        # Use ThreadPoolExecutor for parallelization
-        with ThreadPoolExecutor(max_workers=min(7, multiprocessing.cpu_count())) as executor:
-            # Stage 1: Independent computations that can be parallelized
-            futures = []
-            
-            # Gaussian smoothing
-            def compute_gaussian(sigma):
-                return ndimage.gaussian_filter(image_2d, sigma)
-            
-            for sigma in self.alphas:
-                future = executor.submit(compute_gaussian, sigma)
-                futures.append(('gaussian', sigma, future))
-
-            # Difference of Gaussians
-            def compute_dog(s1, s2):
-                g1 = ndimage.gaussian_filter(image_2d, s1)
-                g2 = ndimage.gaussian_filter(image_2d, s2)
-                return g1 - g2
-            
-            dog_pairs = self.dogs
-            for (s1, s2) in dog_pairs:
-                future = executor.submit(compute_dog, s1, s2)
-                futures.append(('dog', s1, future))
-            
-            # Local statistics computation
-            def compute_local_mean():
-                window_size = self.windows
-                kernel = np.ones((window_size, window_size)) / (window_size**2)
-                return ndimage.convolve(image_2d, kernel, mode='reflect')
-            
-            future = executor.submit(compute_local_mean)
-            futures.append(('local_mean', None, future))
-            
-            def compute_local_variance():
-                window_size = self.windows
-                kernel = np.ones((window_size, window_size)) / (window_size**2)
-                mean = np.mean(image_2d)
-                return ndimage.convolve((image_2d - mean)**2, kernel, mode='reflect')
-            
-            future = executor.submit(compute_local_variance)
-            futures.append(('local_var', None, future))
-            
-            # Gradient computation
-            def compute_gradients():
-                gx = ndimage.sobel(image_2d, axis=1, mode='reflect')  # x direction
-                gy = ndimage.sobel(image_2d, axis=0, mode='reflect')  # y direction
-                return gx, gy
-            
-            future = executor.submit(compute_gradients)
-            futures.append(('gradients', None, future))
-            
-            # Collect results for the independent computations
-            results = {}
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    if task_type == 'gradients':
-                        # Store the gradient components separately
-                        gx, gy = result
-                        results['gx'] = gx
-                        results['gy'] = gy
-                    else:
-                        results[f"{task_type}_{params}" if params is not None else task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
-            
-            # Stage 2: Dependent computations that need results from Stage 1
-            futures = []
-            
-            # Gradient magnitude (depends on gradients)
-            def compute_gradient_magnitude(gx, gy):
-                return np.sqrt(gx**2 + gy**2)
-            
-            future = executor.submit(compute_gradient_magnitude, results['gx'], results['gy'])
-            futures.append(('gradient_magnitude', None, future))
-            
-            # Second-order gradients (depend on first gradients)
-            def compute_second_derivatives(gx, gy):
-                gxx = ndimage.sobel(gx, axis=1, mode='reflect')
-                gyy = ndimage.sobel(gy, axis=0, mode='reflect')
-                # Cross derivatives for Hessian determinant
-                gxy = ndimage.sobel(gx, axis=0, mode='reflect')
-                gyx = ndimage.sobel(gy, axis=1, mode='reflect')
-                return gxx, gyy, gxy, gyx
-            
-            future = executor.submit(compute_second_derivatives, results['gx'], results['gy'])
-            futures.append(('second_derivatives', None, future))
-            
-            # Collect results for the dependent computations
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    if task_type == 'second_derivatives':
-                        # Store the second derivative components separately
-                        gxx, gyy, gxy, gyx = result
-                        results['gxx'] = gxx
-                        results['gyy'] = gyy
-                        results['gxy'] = gxy
-                        results['gyx'] = gyx
-                    else:
-                        results[task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
-            
-            # Stage 3: Final computations that depend on Stage 2 results
-            futures = []
-            
-            # Laplacian and Hessian determinant (depend on second derivatives)
-            def compute_laplacian(gxx, gyy):
-                return gxx + gyy
-            
-            future = executor.submit(compute_laplacian, results['gxx'], results['gyy'])
-            futures.append(('laplacian', None, future))
-            
-            def compute_hessian_det(gxx, gyy, gxy, gyx):
-                return gxx * gyy - gxy * gyx
-            
-            future = executor.submit(compute_hessian_det, 
-                                  results['gxx'], results['gyy'], 
-                                  results['gxy'], results['gyx'])
-            futures.append(('hessian_det', None, future))
-            
-            # Collect final results
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    results[task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type}: {str(e)}")
-        
-        # Organize results in the expected order
-        features = []
-        
-        # Add Gaussian features
-        for sigma in self.alphas:
-            features.append(results[f'gaussian_{sigma}'])
-
-        for sigma in self.dogs:
-            features.append(results[f'dog_{sigma[0]}'])
-        
-        # Add local statistics
-        features.append(results['local_mean'])
-        features.append(results['local_var'])
-        
-        # Add gradient magnitude
-        features.append(results['gradient_magnitude'])
-        
-        # Add Laplacian and Hessian determinant
-        features.append(results['laplacian'])
-        features.append(results['hessian_det'])
-        
-        # Verify shapes
-        for i, feat in enumerate(features):
-            if feat.shape != original_shape:
-                # Check dimensionality and expand if needed
-                if len(feat.shape) < len(original_shape):
-                    feat_adjusted = feat
-                    missing_dims = len(original_shape) - len(feat.shape)
-                    for _ in range(missing_dims):
-                        feat_adjusted = np.expand_dims(feat_adjusted, axis=0)
-                    
-                    if feat_adjusted.shape != original_shape:
-                        raise ValueError(f"Feature {i} has shape {feat.shape}, expected {original_shape}")
-                    
-                    features[i] = feat_adjusted
+        # Pre-compute all Gaussian filters
+        for sigma in all_sigmas:
+            gaussian_cache[sigma] = ndimage.gaussian_filter(image_3d, sigma)
         
-        return np.stack(features, axis=-1)
-
-
-    def compute_feature_maps_cpu_2d_parallel(self, z=None):
-        """Compute feature maps for 2D images using CPU with thread-based parallelism"""
-        image_2d = self.image_3d[z, :, :]
-        original_shape = image_2d.shape
+        # Gaussian smoothing - use cached results
+        for sigma in self.sigmas:
+            features[..., feature_idx] = gaussian_cache[sigma]
+            feature_idx += 1
         
-        # Use ThreadPoolExecutor for parallelization
-        with ThreadPoolExecutor(max_workers=min(7, multiprocessing.cpu_count())) as executor:
-            # Submit tasks for independent computations
-            futures = []
-            
-            # Gaussian smoothing at different scales
-            def compute_gaussian(sigma):
-                return ndimage.gaussian_filter(image_2d, sigma)
-            
-            gaussian_sigmas = self.alphas
-            for sigma in gaussian_sigmas:
-                future = executor.submit(compute_gaussian, sigma)
-                futures.append(('gaussian', sigma, future))
-            
-            # Difference of Gaussians
-            def compute_dog(s1, s2):
-                g1 = ndimage.gaussian_filter(image_2d, s1)
-                g2 = ndimage.gaussian_filter(image_2d, s2)
-                return g1 - g2
-            
-            dog_pairs = self.dogs
-            for (s1, s2) in dog_pairs:
-                future = executor.submit(compute_dog, s1, s2)
-                futures.append(('dog', (s1, s2), future))
-            
-            # Gradient computation
-            def compute_gradient_magnitude():
-                gx = ndimage.sobel(image_2d, axis=1, mode='reflect')  # x direction
-                gy = ndimage.sobel(image_2d, axis=0, mode='reflect')  # y direction
-                return np.sqrt(gx**2 + gy**2)
-            
-            future = executor.submit(compute_gradient_magnitude)
-            futures.append(('gradient_magnitude', None, future))
-            
-            # Collect results
-            results = {}
-            for task_type, params, future in futures:
-                try:
-                    result = future.result()
-                    if params is not None:
-                        if task_type == 'dog':
-                            s1, s2 = params
-                            results[f"{task_type}_{s1}_{s2}"] = result
-                        else:
-                            results[f"{task_type}_{params}"] = result
-                    else:
-                        results[task_type] = result
-                except Exception as e:
-                    raise RuntimeError(f"Error in task {task_type} with params {params}: {str(e)}")
-        
-        # Organize results in the expected order
-        features = []
-        
-        # Add Gaussian features
-        for sigma in gaussian_sigmas:
-            features.append(results[f'gaussian_{sigma}'])
-        
-        # Add Difference of Gaussians features
-        for (s1, s2) in dog_pairs:
-            features.append(results[f'dog_{s1}_{s2}'])
-        
-        # Add gradient magnitude
-        features.append(results['gradient_magnitude'])
-        
-        # Verify shapes
-        for i, feat in enumerate(features):
-            if feat.shape != original_shape:
-                # Check dimensionality and expand if needed
-                if len(feat.shape) < len(original_shape):
-                    feat_adjusted = feat
-                    missing_dims = len(original_shape) - len(feat.shape)
-                    for _ in range(missing_dims):
-                        feat_adjusted = np.expand_dims(feat_adjusted, axis=0)
-                    
-                    if feat_adjusted.shape != original_shape:
-                        raise ValueError(f"Feature {i} has shape {feat.shape}, expected {original_shape}")
-                    
-                    features[i] = feat_adjusted
+        # Difference of Gaussians - use cached results
+        for s1, s2 in self.dogs:
+            features[..., feature_idx] = gaussian_cache[s1] - gaussian_cache[s2]
+            feature_idx += 1
+        
+        # Gaussian gradient magnitudes for each sigma (vectorized)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            gx = ndimage.sobel(gaussian_img, axis=2, mode='reflect')
+            gy = ndimage.sobel(gaussian_img, axis=1, mode='reflect')
+            gz = ndimage.sobel(gaussian_img, axis=0, mode='reflect')
+            features[..., feature_idx] = np.sqrt(gx**2 + gy**2 + gz**2)
+            feature_idx += 1
+        
+        # Laplacian of Gaussian for each sigma (vectorized)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            features[..., feature_idx] = ndimage.laplace(gaussian_img, mode='reflect')
+            feature_idx += 1
+        
+        # Largest Hessian eigenvalue for each sigma (fully vectorized)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            
+            # Compute second derivatives (Hessian components) - all vectorized
+            hxx = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 0, 2], mode='reflect')
+            hyy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 2, 0], mode='reflect')
+            hzz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[2, 0, 0], mode='reflect')
+            hxy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 1, 1], mode='reflect')
+            hxz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[1, 0, 1], mode='reflect')
+            hyz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[1, 1, 0], mode='reflect')
+            
+            # Vectorized eigenvalue computation using numpy broadcasting
+            # Create arrays with shape (d0, d1, d2, 3, 3) for all Hessian matrices
+            shape = image_3d.shape
+            hessian_matrices = np.zeros(shape + (3, 3))
+            
+            # Fill the symmetric Hessian matrices
+            hessian_matrices[..., 0, 0] = hxx
+            hessian_matrices[..., 1, 1] = hyy
+            hessian_matrices[..., 2, 2] = hzz
+            hessian_matrices[..., 0, 1] = hessian_matrices[..., 1, 0] = hxy
+            hessian_matrices[..., 0, 2] = hessian_matrices[..., 2, 0] = hxz
+            hessian_matrices[..., 1, 2] = hessian_matrices[..., 2, 1] = hyz
+            
+            # Reshape for batch eigenvalue computation
+            original_shape = hessian_matrices.shape[:-2]  # (d0, d1, d2)
+            batch_size = np.prod(original_shape)
+            hessian_batch = hessian_matrices.reshape(batch_size, 3, 3)
+            
+            # Compute eigenvalues for all matrices at once
+            eigenvalues_batch = np.real(np.linalg.eigvals(hessian_batch))
+            
+            # Get only the largest eigenvalue for each matrix
+            largest_eigenvalues = np.max(eigenvalues_batch, axis=1)
+            
+            # Reshape back to original spatial dimensions
+            largest_eigenvalues = largest_eigenvalues.reshape(original_shape)
+            
+            # Add the largest eigenvalue as a feature
+            features[..., feature_idx] = largest_eigenvalues
+            feature_idx += 1
         
-        return np.stack(features, axis=-1)
 
+        # Normalize only morphological features, keep intensity features raw
+        intensity_features = features[..., :num_basic_features]  # original + gaussians + DoGs
+        morphology_features = features[..., num_basic_features:]  # gradients + laplacians + eigenvalues
 
-    def compute_feature_maps_cpu_parallel(self, image_3d=None):
-        """Use ThreadPoolExecutor 
+        # Normalize only morphological features
+        morph_means = np.mean(morphology_features, axis=(0,1,2), keepdims=True)
+        morph_stds = np.std(morphology_features, axis=(0,1,2), keepdims=True)
+        morph_stds = np.where(morph_stds == 0, 1, morph_stds)
+        morphology_features = (morphology_features - morph_means) / morph_stds
+
+        # Recombine
+        features = np.concatenate([intensity_features, morphology_features], axis=-1)
         
-        While threads don't give true parallelism for CPU-bound tasks due to the GIL,
-        numpy/scipy release the GIL during computation, so this can still be effective.
-        """
+        return features
+
+    def compute_deep_feature_maps_cpu_smaller(self, image_3d=None): #smaller
+        """Optimized version using determinant instead of full eigenvalue computation. Currently not in use anywhere"""
         if image_3d is None:
             image_3d = self.image_3d
-        original_shape = image_3d.shape
         
-        features = []
+        # Calculate total number of features (using determinant instead of 3 eigenvalues)
+        num_basic_features = 1 + len(self.sigmas) + len(self.dogs)
+        num_gradient_features = len(self.sigmas)
+        num_laplacian_features = len(self.sigmas)
+        num_hessian_features = len(self.sigmas) * 3  # determinant + trace + frobenius norm
         
-        # Using ThreadPoolExecutor which is more compatible with GUI applications
-        with ThreadPoolExecutor(max_workers=min(7, multiprocessing.cpu_count())) as executor:
-            # Submit all tasks to the executor
-            futures = []
-            
-            # Gaussian smoothing at different scales
-            for sigma in self.alphas:
-                future = executor.submit(ndimage.gaussian_filter, image_3d, sigma)
-                futures.append(future)
-            
-            def compute_dog_local(img, s1, s2):
-                g1 = ndimage.gaussian_filter(img, s1) # Consider just having this return the gaussians to
-                g2 = ndimage.gaussian_filter(img, s2)
-                return g1 - g2
+        total_features = num_basic_features + num_gradient_features + num_laplacian_features + num_hessian_features
+        
+        # Pre-allocate result array
+        features = np.empty(image_3d.shape + (total_features,), dtype=image_3d.dtype)
+        features[..., 0] = image_3d
+        
+        feature_idx = 1
+        
+        # Cache for Gaussian filters
+        gaussian_cache = {}
+        all_sigmas = set(self.sigmas)
+        for s1, s2 in self.dogs:
+            all_sigmas.add(s1)
+            all_sigmas.add(s2)
+        
+        # Pre-compute all Gaussian filters
+        for sigma in all_sigmas:
+            gaussian_cache[sigma] = ndimage.gaussian_filter(image_3d, sigma)
+        
+        # Gaussian smoothing
+        for sigma in self.sigmas:
+            features[..., feature_idx] = gaussian_cache[sigma]
+            feature_idx += 1
+        
+        # Difference of Gaussians
+        for s1, s2 in self.dogs:
+            features[..., feature_idx] = gaussian_cache[s1] - gaussian_cache[s2]
+            feature_idx += 1
+        
+        # Gaussian gradient magnitudes
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            gx = ndimage.sobel(gaussian_img, axis=2, mode='reflect')
+            gy = ndimage.sobel(gaussian_img, axis=1, mode='reflect')
+            gz = ndimage.sobel(gaussian_img, axis=0, mode='reflect')
+            features[..., feature_idx] = np.sqrt(gx**2 + gy**2 + gz**2)
+            feature_idx += 1
+        
+        # Laplacian of Gaussian
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            features[..., feature_idx] = ndimage.laplace(gaussian_img, mode='reflect')
+            feature_idx += 1
+        
+        # Hessian-based features (much faster than full eigenvalue computation)
+        for sigma in self.sigmas:
+            gaussian_img = gaussian_cache[sigma]
+            
+            # Compute second derivatives
+            hxx = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 0, 2], mode='reflect')
+            hyy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 2, 0], mode='reflect')
+            hzz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[2, 0, 0], mode='reflect')
+            hxy = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[0, 1, 1], mode='reflect')
+            hxz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[1, 0, 1], mode='reflect')
+            hyz = ndimage.gaussian_filter(gaussian_img, sigma=0, order=[1, 1, 0], mode='reflect')
+            
+            # Hessian determinant (captures overall curvature)
+            determinant = (hxx * (hyy * hzz - hyz**2) - 
+                          hxy * (hxy * hzz - hxz * hyz) + 
+                          hxz * (hxy * hyz - hyy * hxz))
+            features[..., feature_idx] = determinant
+            feature_idx += 1
+            
+            # Hessian trace (sum of eigenvalues)
+            trace = hxx + hyy + hzz
+            features[..., feature_idx] = trace
+            feature_idx += 1
+            
+            # Frobenius norm (overall curvature magnitude)
+            frobenius_norm = np.sqrt(hxx**2 + hyy**2 + hzz**2 + 2*(hxy**2 + hxz**2 + hyz**2))
+            features[..., feature_idx] = frobenius_norm
+            feature_idx += 1
 
-            # Difference of Gaussians
-            for (s1, s2) in self.dogs:
-                
-                future = executor.submit(compute_dog_local, image_3d, s1, s2)
-                futures.append(future)
-            
-            # Gradient magnitude
-            def compute_gradient_local(img):
-                gx = ndimage.sobel(img, axis=2, mode='reflect')
-                gy = ndimage.sobel(img, axis=1, mode='reflect')
-                gz = ndimage.sobel(img, axis=0, mode='reflect')
-                return np.sqrt(gx**2 + gy**2 + gz**2)
-            
-            future = executor.submit(compute_gradient_local, image_3d)
-            futures.append(future)
+        """
+        # Normalize features: zero-mean, unit variance per feature band
+        # Compute mean and std across spatial dimensions (0,1,2), keeping feature dimension
+        feature_means = np.mean(features, axis=(0, 1, 2), keepdims=True)
+        feature_stds = np.std(features, axis=(0, 1, 2), keepdims=True)
+        
+        # Avoid division by zero for constant features
+        feature_stds = np.where(feature_stds == 0, 1, feature_stds)
+        
+        # Normalize in-place for memory efficiency
+        features = (features - feature_means) / feature_stds
+        """
+        # Normalize only morphological features, keep intensity features raw
+        intensity_features = features[..., :num_basic_features]  # original + gaussians + DoGs
+        morphology_features = features[..., num_basic_features:]  # gradients + laplacians + eigenvalues
+
+        # Normalize only morphological features
+        morph_means = np.mean(morphology_features, axis=(0,1,2), keepdims=True)
+        morph_stds = np.std(morphology_features, axis=(0,1,2), keepdims=True)
+        morph_stds = np.where(morph_stds == 0, 1, morph_stds)
+        morphology_features = (morphology_features - morph_means) / morph_stds
+
+        # Recombine
+        features = np.concatenate([intensity_features, morphology_features], axis=-1)
+
+        return features
+
+
+    def compute_feature_maps_cpu(self, image_3d=None): #lil
+        """Optimized version that caches Gaussian filters to avoid redundant computation"""
+        if image_3d is None:
+            image_3d = self.image_3d
+
+        if image_3d.ndim == 4 and image_3d.shape[-1] == 3:
+            # RGB case - process each channel
+            features_per_channel = []
+            for channel in range(3):
+                channel_features = self.compute_feature_maps_cpu(image_3d[..., channel])
+                features_per_channel.append(channel_features)
             
-            # Collect results
-            for future in futures:
-                result = future.result()
-                features.append(result)
+            # Stack all channel features
+            return np.concatenate(features_per_channel, axis=-1)
+        
+        # Pre-allocate result array
+        num_features = len(self.sigmas) + len(self.dogs) + 2
+        features = np.empty(image_3d.shape + (num_features,), dtype=image_3d.dtype)
+        features[..., 0] = image_3d
+        
+        feature_idx = 1
+        
+        # Cache for Gaussian filters - only compute each sigma once
+        gaussian_cache = {}
         
-        # Verify shapes
-        for i, feat in enumerate(features):
-            if feat.shape != original_shape:
-                feat_adjusted = np.expand_dims(feat, axis=0)
-                if feat_adjusted.shape != original_shape:
-                    raise ValueError(f"Feature {i} has shape {feat.shape}, expected {original_shape}")
-                features[i] = feat_adjusted
+        # Compute all unique sigmas needed (from both sigmas and dogs)
+        all_sigmas = set(self.sigmas)
+        for s1, s2 in self.dogs:
+            all_sigmas.add(s1)
+            all_sigmas.add(s2)
         
-        return np.stack(features, axis=-1)
+        # Pre-compute all Gaussian filters
+        for sigma in all_sigmas:
+            gaussian_cache[sigma] = ndimage.gaussian_filter(image_3d, sigma)
+        
+        # Gaussian smoothing - use cached results
+        for sigma in self.sigmas:
+            features[..., feature_idx] = gaussian_cache[sigma]
+            feature_idx += 1
+        
+        # Difference of Gaussians - use cached results
+        for s1, s2 in self.dogs:
+            features[..., feature_idx] = gaussian_cache[s1] - gaussian_cache[s2]
+            feature_idx += 1
+        
+        # Gradient magnitude
+        gx = ndimage.sobel(image_3d, axis=2, mode='reflect')
+        gy = ndimage.sobel(image_3d, axis=1, mode='reflect')
+        gz = ndimage.sobel(image_3d, axis=0, mode='reflect')
+        features[..., feature_idx] = np.sqrt(gx**2 + gy**2 + gz**2)
 
+        return features
 
     def organize_by_z(self, coordinates):
         """
@@ -692,113 +549,71 @@ class InteractiveSegmenter:
 
     def process_chunk(self, chunk_coords):
         """
-        Process a chunk of coordinates, handling both mem_lock and non-mem_lock cases.
-        Uses a consistent approach based on coordinates.
+        Vectorized process_chunk that releases GIL more effectively
+        """
+        if self.realtimechunks is None:
+            # Generate coordinates using vectorized operations
+            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()
+            ])
+        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]
         
-        Parameters:
-        -----------
-        chunk_coords : list of tuples
-            List of (z,y,x) coordinate tuples to process
+        # 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)
         
-        Returns:
-        --------
-        tuple : (foreground, background)
-            Sets of coordinates classified as foreground or background
-        """
-        foreground = set()
-        background = set()
+        # 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])
         
-        if self.previewing or not self.use_two:
-
-            if self.realtimechunks is None: #Presuming we're segmenting all
-                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]
-
-                # Consider moving this to process chunk ??
-                chunk_coords = np.stack(np.meshgrid(
-                    np.arange(z_min, z_max),
-                    np.arange(y_min, y_max),
-                    np.arange(x_min, x_max),
-                    indexing='ij'
-                )).reshape(3, -1).T
-                
-                chunk_coords = (list(map(tuple, chunk_coords)))
-            else: #Presumes we're not segmenting all
-                # Find min/max bounds of the coordinates to get the smallest containing subarray
-                z_coords = [z for z, y, x in chunk_coords]
-                y_coords = [y for z, y, x in chunk_coords]
-                x_coords = [x for z, y, x in chunk_coords]
-                
-                z_min, z_max = min(z_coords), max(z_coords)
-                y_min, y_max = min(y_coords), max(y_coords)
-                x_min, x_max = min(x_coords), max(x_coords)
-
-            
-            # Extract the subarray
-            subarray = self.image_3d[z_min:z_max+1, y_min:y_max+1, x_min:x_max+1]
-            
-            # Compute features for this subarray
-            if self.speed:
-                feature_map = self.compute_feature_maps_cpu_parallel(subarray) #If the interactive segmenter is slow
-            else:                                                              #Due to the parallel, consider singleton implementation for it specifically
-                feature_map = self.compute_deep_feature_maps_cpu_parallel(subarray)
-            
-            # Extract features for each coordinate, adjusting for subarray offset
-            features = []
-            for z, y, x in chunk_coords:
-                # Transform global coordinates to local subarray coordinates
-                local_z = z - z_min
-                local_y = y - y_min
-                local_x = x - x_min
-                
-                # Get feature at this position
-                feature = feature_map[local_z, local_y, local_x]
-                features.append(feature)
-                    
-            
-            # Make predictions
+        # 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
             predictions = self.model.predict(features)
-            
-            # Assign coordinates based on predictions
-            for coord, pred in zip(chunk_coords, predictions):
-                if pred:
-                    foreground.add(coord)
-                else:
-                    background.add(coord)
-
         else:
-
-            chunk_coords = self.twodim_coords(chunk_coords[0], chunk_coords[1], chunk_coords[2], chunk_coords[3], chunk_coords[4])
-
-            chunk_coords = self.organize_by_z(chunk_coords)
-
-            for z, coords in chunk_coords.items():
-
-                if self.feature_cache is None:
-                    features = self.get_feature_map_slice(z, self.speed, self.cur_gpu)
-                    features = [features[y, x] for y, x in coords]
-                elif z not in self.feature_cache and not self.previewing:
-                    features = self.get_feature_map_slice(z, self.speed, self.cur_gpu)
-                    features = [features[y, x] for y, x in coords]
-                elif z not in self.feature_cache or self.feature_cache is None and self.previewing:
-                    features = self.map_slice
-                    try:
-                        features = [features[y, x] for y, x in coords]
-                    except:
-                        return [], []
-                else:  
-                    features = [self.feature_cache[z][y, x] for y, x in coords]
-
-                predictions = self.model.predict(features)
-                
-                for (y, x), pred in zip(coords, predictions):
-                    coord = (z, y, x)  # Reconstruct the 3D coordinate as a tuple
-                    if pred:
-                        foreground.add(coord)
-                    else:
-                        background.add(coord)
+            # 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
 
     def twodim_coords(self, y_dim, x_dim, z, chunk_size = None, subrange = None):
@@ -861,62 +676,38 @@ class InteractiveSegmenter:
 
 
     def segment_volume(self, array, chunk_size=None, gpu=False):
-        """Segment volume using parallel processing of chunks with vectorized chunk creation"""
-        #Change the above chunk size to None to have it auto-compute largest chunks (not sure which is faster, 64 seems reasonable in test cases)
-
-        self.realtimechunks = None # Presumably no longer need this.
-        self.map_slice = None
-
-        chunk_size = self.master_chunk #memory efficient chunk
-
+        """
+        Process chunks in batches equal to CPU cores for optimal GIL handling
+        """
+        
+        self.realtimechunks = None
+        chunk_size = self.master_chunk
 
         def create_2d_chunks():
-            """
-            Create chunks by z-slices for 2D processing.
-            Each chunk is a complete z-slice with all y,x coordinates,
-            unless the slice exceeds 262144 pixels, in which case it's divided into subchunks.
-            
-            Returns:
-                List of chunks, where each chunk contains the coordinates for one z-slice or subchunk
-            """
-            MAX_CHUNK_SIZE = 262144
-
+            """Same as your existing implementation"""
+            MAX_CHUNK_SIZE = self.twod_chunk_size
             chunks = []
             
             for z in range(self.image_3d.shape[0]):
-                # Get the dimensions of this z-slice
                 y_dim = self.image_3d.shape[1]
                 x_dim = self.image_3d.shape[2]
                 total_pixels = y_dim * x_dim
                 
-                # If the slice is small enough, do not subchunk
                 if total_pixels <= MAX_CHUNK_SIZE:
-
                     chunks.append([y_dim, x_dim, z, total_pixels, None])
-
                 else:
-                    # Determine which dimension to divide (the largest one)
                     largest_dim = 'y' if y_dim >= x_dim else 'x'
-                    
-                    # Calculate how many divisions we need
                     num_divisions = int(np.ceil(total_pixels / MAX_CHUNK_SIZE))
                     
-                    # Calculate the approx size of each division along the largest dimension
                     if largest_dim == 'y':
                         div_size = int(np.ceil(y_dim / num_divisions))
-                        # Create subchunks by dividing the y-dimension
                         for i in range(0, y_dim, div_size):
                             end_i = min(i + div_size, y_dim)
-
-
                             chunks.append([y_dim, x_dim, z, None, ['y', i, end_i]])
-
-                    else:  # largest_dim == 'x'
+                    else:
                         div_size = int(np.ceil(x_dim / num_divisions))
-                        # Create subchunks by dividing the x-dimension
                         for i in range(0, x_dim, div_size):
                             end_i = min(i + div_size, x_dim)
-
                             chunks.append([y_dim, x_dim, z, None, ['x', i, end_i]])
             
             return chunks
@@ -924,30 +715,20 @@ class InteractiveSegmenter:
         print("Chunking data...")
         
         if not self.use_two:
-            # Determine optimal chunk size based on number of cores if not specified
+            # Create smaller chunks for better load balancing
             if chunk_size is None:
                 total_cores = multiprocessing.cpu_count()
-                
-                # Calculate total volume and target volume per core
                 total_volume = np.prod(self.image_3d.shape)
-                target_volume_per_chunk = total_volume / total_cores
+                target_volume_per_chunk = total_volume / (total_cores * 4)  # 4x more chunks
                 
-                # Calculate chunk size that would give us roughly one chunk per core
-                # Using cube root since we want roughly equal sizes in all dimensions
                 chunk_size = int(np.cbrt(target_volume_per_chunk))
-                
-                # Ensure chunk size is at least 32 (minimum reasonable size) and not larger than smallest dimension
-                chunk_size = max(32, min(chunk_size, min(self.image_3d.shape)))
-                
-                # Round to nearest multiple of 32 for better memory alignment
-                chunk_size = ((chunk_size + 15) // 32) * 32
+                chunk_size = max(16, min(chunk_size, min(self.image_3d.shape) // 2))
+                chunk_size = ((chunk_size + 7) // 16) * 16
             
-            # Calculate number of chunks in each dimension
             z_chunks = (self.image_3d.shape[0] + chunk_size - 1) // chunk_size
             y_chunks = (self.image_3d.shape[1] + chunk_size - 1) // chunk_size
             x_chunks = (self.image_3d.shape[2] + chunk_size - 1) // chunk_size
             
-            # Create start indices for all chunks at once
             chunk_starts = np.array(np.meshgrid(
                 np.arange(z_chunks) * chunk_size,
                 np.arange(y_chunks) * chunk_size,
@@ -960,38 +741,148 @@ class InteractiveSegmenter:
                 z_end = min(z_start + chunk_size, self.image_3d.shape[0])
                 y_end = min(y_start + chunk_size, self.image_3d.shape[1])
                 x_end = min(x_start + chunk_size, self.image_3d.shape[2])
-                
                 coords = [z_start, z_end, y_start, y_end, x_start, x_end]
                 chunks.append(coords)
-
-
         else:
             chunks = create_2d_chunks()
-            self.feature_cache = None #Decided this should not maintain training data for segmenting 2D
-        
-        foreground_coords = set()
-        background_coords = set()
-
-        print("Segmenting chunks...")
-
-        for i, chunk in enumerate(chunks):
-            fore, _ = self.process_chunk(chunk)
-            fg_array = np.array(list(fore))
-            del fore
-            if len(fg_array) > 0:  # Check if we have any foreground coordinates
-                # Unpack into separate coordinate arrays
-                z_coords, y_coords, x_coords = fg_array[:, 0], fg_array[:, 1], fg_array[:, 2]
-                # Assign values in a single vectorized operation
-                array[z_coords, y_coords, x_coords] = 255
-            try:
-                chunk[i] = None #Help garbage collection
-            except:
-                pass
-            print(f"Processed {i}/{len(chunks)} chunks")
+
+        print("Processing chunks in batches...")
+        
+        # Process chunks in batches equal to CPU count
+        max_workers = self.batch_amplifier * multiprocessing.cpu_count()
+        batch_size = max_workers  # One batch per core
+        total_processed = 0
+        
+        # Configure sklearn for maximum parallelism
+        if hasattr(self.model, 'n_jobs'):
+            original_n_jobs = self.model.n_jobs
+            self.model.n_jobs = -1  # Use all cores for sklearn prediction
+        
+        try:
+            for batch_start in range(0, len(chunks), batch_size):
+                batch_end = min(batch_start + batch_size, len(chunks))
+                chunk_batch = chunks[batch_start:batch_end]
+                
+                print(f"Processing batch {batch_start//batch_size + 1}/{(len(chunks) + batch_size - 1)//batch_size}")
+                
+                # PHASE 1: Extract features in parallel (releases GIL)
+                batch_results = []
+                
+                with ThreadPoolExecutor(max_workers=len(chunk_batch)) as executor:
+                    futures = []
+                    for chunk in chunk_batch:
+                        future = executor.submit(self.extract_chunk_features, chunk)
+                        futures.append(future)
+                    
+                    # Collect feature results
+                    for future in futures:
+                        features, coords = future.result()
+                        if len(features) > 0:
+                            batch_results.append((features, coords))
+                
+                # PHASE 2: Batch predict with sklearn's parallelism (no GIL issues)
+                if batch_results:
+                    # Combine all features from this batch
+                    all_batch_features = np.vstack([result[0] for result in batch_results])
+                    all_batch_coords = np.vstack([result[1] for result in batch_results])
+                    
+                    # Single prediction call using sklearn's internal parallelism
+                    predictions = self.model.predict(all_batch_features)
+                    predictions = np.array(predictions, dtype=bool)
+                    
+                    # Apply predictions to array
+                    foreground_coords = all_batch_coords[predictions]
+                    if len(foreground_coords) > 0:
+                        z_coords, y_coords, x_coords = foreground_coords[:, 0], foreground_coords[:, 1], foreground_coords[:, 2]
+                        array[z_coords, y_coords, x_coords] = 255
+                    
+                    # Clean up batch data for memory management
+                    del all_batch_features, all_batch_coords, predictions, foreground_coords
+                    
+                total_processed += len(chunk_batch)
+                print(f"Completed {total_processed}/{len(chunks)} chunks")
+        
+        finally:
+            # Restore original sklearn settings
+            if hasattr(self.model, 'n_jobs'):
+                self.model.n_jobs = original_n_jobs
         
-        #Ok so this should be returned one chunk at a time I presume.
         return array
 
+    def extract_chunk_features(self, chunk_coords):
+        """
+        Extract features for a single chunk without prediction
+        Designed to release GIL effectively
+        """
+        
+        if self.previewing or not self.use_two:
+            if self.realtimechunks is None:
+                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]
+
+                # Vectorized coordinate generation (releases GIL)
+                z_range = np.arange(z_min, z_max)
+                y_range = np.arange(y_min, y_max)
+                x_range = np.arange(x_min, x_max)
+                
+                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()
+                ])
+            else:
+                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 and compute features (releases GIL)
+            subarray = self.image_3d[z_min:z_max+1, y_min:y_max+1, x_min:x_max+1]
+            
+            if self.speed:
+                feature_map = self.compute_feature_maps_cpu(subarray)
+            else:
+                feature_map = self.compute_deep_feature_maps_cpu(subarray)
+            
+            # Vectorized feature extraction (releases GIL)
+            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]]
+            
+            return features, chunk_coords_array
+        
+        else:
+            # Handle 2D case
+            chunk_coords_list = self.twodim_coords(chunk_coords[0], chunk_coords[1], 
+                                                 chunk_coords[2], chunk_coords[3], chunk_coords[4])
+            chunk_coords_by_z = self.organize_by_z(chunk_coords_list)
+            
+            all_features = []
+            all_coords = []
+            
+            for z, coords in chunk_coords_by_z.items():
+                coords_array = np.array(coords)
+                
+                # Get features for this z-slice
+                features_slice = self.get_feature_map_slice(z, self.speed, self.cur_gpu)
+                features = features_slice[coords_array[:, 0], coords_array[:, 1]]
+
+                
+                # Convert to 3D coordinates
+                coords_3d = np.column_stack([
+                    np.full(len(coords_array), z),
+                    coords_array[:, 0],
+                    coords_array[:, 1]
+                ])
+                
+                all_features.append(features)
+                all_coords.append(coords_3d)
+            
+            if all_features:
+                return np.vstack(all_features), np.vstack(all_coords)
+            else:
+                return np.array([]), np.array([])
+                
     def update_position(self, z=None, x=None, y=None):
         """Update current position for chunk prioritization with safeguards"""
         
@@ -1019,10 +910,7 @@ class InteractiveSegmenter:
         
         # Only clear map_slice if z changes and we're not already generating a new one
         if self.current_z != self.prev_z:
-            # Instead of setting to None, check if we already have it in the cache
-            if hasattr(self, 'feature_cache') and self.feature_cache is not None:
-                if self.current_z not in self.feature_cache:
-                    self.map_slice = None
+
             self._currently_segmenting = None
         
         # Update previous z
@@ -1077,12 +965,115 @@ class InteractiveSegmenter:
 
         print("Ready!")
 
+    def get_realtime_chunks_2d(self, chunk_size=None):
+        """
+        Create square chunks with 1 z-thickness (2D chunks across XY planes)
+        """
+        
+        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
+        
+        # Populate chunk dictionary
+        chunk_dict = {}
+        
+        # Create chunks for each Z plane (single Z thickness)
+        for z in range(self.image_3d.shape[0]):
+            if small_dims:
+                chunk_dict[(z, 0, 0)] = {
+                    'coords': [0, self.image_3d.shape[1], 0, self.image_3d.shape[2]],
+                    'processed': False,
+                    'z': z  # Keep for backward compatibility
+                }
+            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],
+                            'processed': False,
+                            'z': z  # Keep for backward compatibility
+                        }
+        
+        self.realtimechunks = chunk_dict
+        print("Ready!")
+
+    def process_slice_features(self, z: 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]]:
+        """
+        Helper function to process a single slice and extract features.
+        Returns tuple of (foreground_features, background_features) for this slice.
+        """
+        slice_foreground_features = []
+        slice_background_features = []
+        
+        current_map = self.get_feature_map_slice(z, speed, use_gpu)
+        
+        if z in z_fores:
+            for y, x in z_fores[z]:
+                feature_vector = current_map[y, x]
+                slice_foreground_features.append(feature_vector)
+        
+        if z in z_backs:
+            for y, x in z_backs[z]:
+                feature_vector = current_map[y, x]
+                slice_background_features.append(feature_vector)
+        
+        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]]:
+        """
+        Process feature extraction using ThreadPoolExecutor for parallel execution.
+        """
+        max_cores = multiprocessing.cpu_count()
+        foreground_features = []
+        background_features = []
+        
+        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
+            }
+            
+            # 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)
+        
+        return foreground_features, background_features
 
     def segment_volume_realtime(self, gpu = False):
 
 
         if self.realtimechunks is None:
-            self.get_realtime_chunks()
+            if not self.use_two:
+                self.get_realtime_chunks()
+            else:
+                self.get_realtime_chunks_2d()
         else:
             for chunk_pos in self.realtimechunks:  # chunk_pos is the (z, y_start, x_start) tuple
                 self.realtimechunks[chunk_pos]['processed'] = False
@@ -1170,45 +1161,121 @@ class InteractiveSegmenter:
             except:
                 pass
 
+    def process_grid_cell(self, grid_cell_info):
+        """
+        Process a single grid cell and return foreground and background features.
+        
+        Args:
+            grid_cell_info: tuple of (grid_z, grid_y, grid_x, box_size, depth, height, width, foreground_array)
+        
+        Returns:
+            tuple: (foreground_features, background_features)
+        """
+        grid_z, grid_y, grid_x, box_size, depth, height, width, foreground_array = grid_cell_info
+        
+        # Calculate the boundaries of this grid cell
+        z_min = grid_z * box_size
+        y_min = grid_y * box_size
+        x_min = grid_x * box_size
+        
+        z_max = min(z_min + box_size, depth)
+        y_max = min(y_min + box_size, height)
+        x_max = min(x_min + box_size, width)
+        
+        # Extract the subarray
+        subarray = self.image_3d[z_min:z_max, y_min:y_max, x_min:x_max]
+        subarray2 = foreground_array[z_min:z_max, y_min:y_max, x_min:x_max]
+        
+        # Compute features for this subarray
+        if self.speed:
+            subarray_features = self.compute_feature_maps_cpu(subarray)
+        else:
+            subarray_features = self.compute_deep_feature_maps_cpu(subarray)
+        
+        # Extract foreground features
+        local_fore_coords = np.argwhere(subarray2 == 1)
+        foreground_features = []
+        for local_z, local_y, local_x in local_fore_coords:
+            feature = subarray_features[local_z, local_y, local_x]
+            foreground_features.append(feature)
+        
+        # Extract background features
+        local_back_coords = np.argwhere(subarray2 == 2)
+        background_features = []
+        for local_z, local_y, local_x in local_back_coords:
+            feature = subarray_features[local_z, local_y, local_x]
+            background_features.append(feature)
+        
+        return foreground_features, background_features
+
+    # Modified main processing code
+    def process_grid_cells_parallel(self, grid_cells_with_scribbles, box_size, depth, height, width, foreground_array, max_workers=None):
+        """
+        Process grid cells in parallel using ThreadPoolExecutor.
+        
+        Args:
+            grid_cells_with_scribbles: List of grid cell coordinates
+            box_size: Size of each grid cell
+            depth, height, width: Dimensions of the 3D image
+            foreground_array: Array marking foreground/background points
+            max_workers: Maximum number of threads (None for default)
+        
+        Returns:
+            tuple: (foreground_features, background_features)
+        """
+        # Prepare data for each grid cell
+        grid_cell_data = [
+            (grid_z, grid_y, grid_x, box_size, depth, height, width, foreground_array)
+            for grid_z, grid_y, grid_x in grid_cells_with_scribbles
+        ]
+        
+        foreground_features = []
+        background_features = []
+        
+        # Process grid cells in parallel
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            # Submit all tasks
+            futures = [executor.submit(self.process_grid_cell, cell_data) for cell_data in grid_cell_data]
+            
+            # Collect results as they complete
+            for future in futures:
+                fore_features, back_features = future.result()
+                foreground_features.extend(fore_features)
+                background_features.extend(back_features)
+        
+        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"""
 
         if not saving:
             print("Training model...")
+            self.model = RandomForestClassifier(
+                n_estimators=100,
+                n_jobs=-1,
+                max_depth=None
+            )
+
         self.speed = speed
         self.cur_gpu = use_gpu
-        if mem_lock != self.mem_lock:
-            self.realtimechunks = None #dump ram
-            self.feature_cache = None
+
+        if use_two != self.use_two:
+            self.realtimechunks = None
 
         if not use_two:
             self.use_two = False
 
         self.mem_lock = mem_lock
 
-        if self.current_speed != speed:
-            self.feature_cache = None
-
-            self.model = RandomForestClassifier(
-                n_estimators=100,
-                n_jobs=-1,
-                max_depth=None
-            )
-
-
         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
-                self.feature_cache = None
                 self.use_two = True
 
-            self.feature_cache = None #Decided this should reset, can remove this line to have it retain prev feature maps
             self.two_slices = []
 
-            if self.feature_cache == None:
-                self.feature_cache = {}
 
             # Get foreground coordinates and features
             z_fore, y_fore, x_fore = np.where(foreground_array == 1)
@@ -1221,14 +1288,6 @@ class InteractiveSegmenter:
 
             back_coords = list(zip(z_back, y_back, x_back))
 
-
-            #slices = set(list(z_back) + list(z_fore))
-
-            #for z in slices:
-                #if z not in self.two_slices:
-                    #changed.append(z)
-                    #self.two_slices.append(z) #Tracks assigning coords to feature map slices
-
             foreground_features = []
             background_features = []
 
@@ -1236,28 +1295,9 @@ class InteractiveSegmenter:
             z_backs = self.organize_by_z(back_coords)
             slices = set(list(z_fores.keys()) + list(z_backs.keys()))
 
-            for z in slices:
-
-
-                current_map = self.get_feature_map_slice(z, speed, use_gpu)
-
-                if z in z_fores:
-                
-                    for y, x in z_fores[z]:
-                        # Get the feature vector for this foreground point
-                        feature_vector = current_map[y, x]
-                        
-                        # Add to our collection
-                        foreground_features.append(feature_vector)
-
-                if z in z_backs:
-                
-                    for y, x in z_backs[z]:
-                        # Get the feature vector for this foreground point
-                        feature_vector = current_map[y, x]
-                    
-                        # Add to our collection
-                        background_features.append(feature_vector)
+            foreground_features, background_features = self.extract_features_parallel(
+                slices, speed, use_gpu, z_fores, z_backs
+            )
 
 
         else: #Forces ram efficiency
@@ -1304,49 +1344,8 @@ class InteractiveSegmenter:
             coord_mapping = {}
             
             # Step 2: Process each grid cell that contains scribbles
-            for grid_z, grid_y, grid_x in grid_cells_with_scribbles:
-                # Calculate the boundaries of this grid cell
-                z_min = grid_z * box_size
-                y_min = grid_y * box_size
-                x_min = grid_x * box_size
-                
-                z_max = min(z_min + box_size, depth)
-                y_max = min(y_min + box_size, height)
-                x_max = min(x_min + box_size, width)
-                
-                # Extract the subarray
-                subarray = self.image_3d[z_min:z_max, y_min:y_max, x_min:x_max]
-                subarray2 = foreground_array[z_min:z_max, y_min:y_max, x_min:x_max]
-                
-                # Compute features for this subarray
-                if self.speed:
-                    subarray_features = self.compute_feature_maps_cpu_parallel(subarray)
-                else:
-                    subarray_features = self.compute_deep_feature_maps_cpu_parallel(subarray)
-                
-                # For each foreground point in this grid cell, extract its feature
-                # Extract foreground features using a direct mask comparison
-                local_fore_coords = np.argwhere(subarray2 == 1)
-                for local_z, local_y, local_x in local_fore_coords:
-                    feature = subarray_features[local_z, local_y, local_x]
-                    foreground_features.append(feature)
-                
-                # Extract background features using a direct mask comparison
-                local_back_coords = np.argwhere(subarray2 == 2)
-                for local_z, local_y, local_x in local_back_coords:
-                    feature = subarray_features[local_z, local_y, local_x]
-                    background_features.append(feature)
-            try:
-                # Get foreground coordinates and features
-                z_fore, y_fore, x_fore = np.where(foreground_array == 1)
-                foreground_features = self.feature_cache[z_fore, y_fore, x_fore]
-
-                # Get background coordinates and features
-                z_back, y_back, x_back = np.where(foreground_array == 2)
-                background_features = self.feature_cache[z_back, y_back, x_back]
-            except:
-                pass
 
+            foreground_features, background_features = self.process_grid_cells_parallel(grid_cells_with_scribbles, box_size, depth, height, width, foreground_array)
 
         if self.previous_foreground is not None:
             failed = True
@@ -1378,7 +1377,6 @@ class InteractiveSegmenter:
         # Combine features and labels
         X = np.vstack([foreground_features, background_features])
         y = np.hstack([np.ones(len(z_fore)), np.zeros(len(z_back))])
-
         
         # Train the model
         try:
@@ -1412,7 +1410,7 @@ class InteractiveSegmenter:
                  use_two=self.use_two,
                  mem_lock=self.mem_lock)
 
-        print(f"Model data saved to {file_name}. Please retrain current model prior to segmentation.")
+        print(f"Model data saved to {file_name}.")
 
 
     def load_model(self, file_name):
@@ -1448,10 +1446,10 @@ class InteractiveSegmenter:
             return
 
         if speed:
-            output = self.compute_feature_maps_cpu_2d_parallel(z = z)
+            output = self.compute_feature_maps_cpu_2d(z = z)
 
         elif not speed:
-            output = self.compute_deep_feature_maps_cpu_2d_parallel(z = z)
+            output = self.compute_deep_feature_maps_cpu_2d(z = z)
 
         return output