nettracer3d 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

nettracer3d/segmenter.py

@@ -20,7 +20,6 @@ from scipy import ndimage
 import multiprocessing
 from collections import defaultdict
 
-
 class InteractiveSegmenter:
     def __init__(self, image_3d, use_gpu=True):
         self.image_3d = image_3d
@@ -81,6 +80,7 @@ class InteractiveSegmenter:
         self._currently_processing = False
         self._skip_next_update = False
         self._last_processed_slice = None
+        self.mem_lock = False
 
     def segment_slice_chunked(self, slice_z, block_size=64):
         """
@@ -203,42 +203,44 @@ class InteractiveSegmenter:
             if self._currently_processing == slice_z:
                 self._currently_processing = None
 
-    def compute_deep_feature_maps_cpu(self):
+    def compute_deep_feature_maps_cpu(self, image_3d = None):
         """Compute feature maps using CPU"""
         features = []
-        original_shape = self.image_3d.shape
+        if image_3d is None:
+            image_3d = self.image_3d
+        original_shape = image_3d.shape
         
         # Gaussian and DoG using scipy
-        print("Obtaining gaussians")
+        #print("Obtaining gaussians")
         for sigma in [0.5, 1.0, 2.0, 4.0]:
-            smooth = ndimage.gaussian_filter(self.image_3d, sigma)
+            smooth = ndimage.gaussian_filter(image_3d, sigma)
             features.append(smooth)
         
-        print("Computing local statistics")
+        #print("Computing local statistics")
         # Local statistics using scipy's convolve
         window_size = 5
         kernel = np.ones((window_size, window_size, window_size)) / (window_size**3)
         
         # Local mean
-        local_mean = ndimage.convolve(self.image_3d, kernel, mode='reflect')
+        local_mean = ndimage.convolve(image_3d, kernel, mode='reflect')
         features.append(local_mean)
         
         # Local variance
-        mean = np.mean(self.image_3d)
-        local_var = ndimage.convolve((self.image_3d - mean)**2, kernel, mode='reflect')
+        mean = np.mean(image_3d)
+        local_var = ndimage.convolve((image_3d - mean)**2, kernel, mode='reflect')
         features.append(local_var)
         
-        print("Computing sobel and gradients")
+        #print("Computing sobel and gradients")
         # Gradient computations using scipy
-        gx = ndimage.sobel(self.image_3d, axis=2, mode='reflect')
-        gy = ndimage.sobel(self.image_3d, axis=1, mode='reflect')
-        gz = ndimage.sobel(self.image_3d, axis=0, mode='reflect')
+        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')
         
         # Gradient magnitude
         gradient_magnitude = np.sqrt(gx**2 + gy**2 + gz**2)
         features.append(gradient_magnitude)
         
-        print("Computing second-order features")
+        #print("Computing second-order features")
         # Second-order gradients
         gxx = ndimage.sobel(gx, axis=2, mode='reflect')
         gyy = ndimage.sobel(gy, axis=1, mode='reflect')
@@ -252,7 +254,162 @@ class InteractiveSegmenter:
         hessian_det = gxx * gyy * gzz
         features.append(hessian_det)
         
-        print("Verifying shapes")
+        #print("Verifying 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
+        
+        return np.stack(features, axis=-1)
+
+    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
+        
+        original_shape = image_3d.shape
+        
+        # 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 [0.5, 1.0, 2.0, 4.0]:
+                future = executor.submit(compute_gaussian, sigma)
+                futures.append(('gaussian', sigma, future))
+            
+            # Local statistics computation
+            def compute_local_mean():
+                window_size = 5
+                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 = 5
+                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
+            
+            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)}")
+        
+        # Organize results in the expected order
+        features = []
+        
+        # Add Gaussian features
+        for sigma in [0.5, 1.0, 2.0, 4.0]:
+            features.append(results[f'gaussian_{sigma}'])
+        
+        # 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:
                 feat_adjusted = np.expand_dims(feat, axis=0)
@@ -323,6 +480,165 @@ class InteractiveSegmenter:
         
         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
+        
+        # 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 [0.5, 1.0, 2.0, 4.0]:
+                future = executor.submit(compute_gaussian, sigma)
+                futures.append(('gaussian', sigma, future))
+            
+            # Local statistics computation
+            def compute_local_mean():
+                window_size = 5
+                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 = 5
+                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 [0.5, 1.0, 2.0, 4.0]:
+            features.append(results[f'gaussian_{sigma}'])
+        
+        # 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
+        
+        return np.stack(features, axis=-1)
+
     def compute_feature_maps(self):
         """Compute all feature maps using GPU acceleration"""
         #if not self.use_gpu:
@@ -513,38 +829,186 @@ class InteractiveSegmenter:
         
         return np.stack(features, axis=-1)
 
-    def compute_feature_maps_cpu(self):
+    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
+        
+        # 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 = [0.5, 1.0, 2.0, 4.0]
+            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 = [(1, 2), (2, 4)]
+            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
+        
+        return np.stack(features, axis=-1)
+
+    def compute_feature_maps_cpu(self, image_3d = None):
         """Compute feature maps using CPU"""
         features = []
-        original_shape = self.image_3d.shape
+        if image_3d is None:
+            image_3d = self.image_3d
+
+        original_shape = image_3d.shape
+
         
         # Gaussian smoothing at different scales
-        print("Obtaining gaussians")
+        #print("Obtaining gaussians")
         for sigma in [0.5, 1.0, 2.0, 4.0]:
-            smooth = ndimage.gaussian_filter(self.image_3d, sigma)
+            smooth = ndimage.gaussian_filter(image_3d, sigma)
             features.append(smooth)
         
-        print("Obtaining dif of gaussians")
+        #print("Obtaining dif of gaussians")
         # Difference of Gaussians
         for (s1, s2) in [(1, 2), (2, 4)]:
-            g1 = ndimage.gaussian_filter(self.image_3d, s1)
-            g2 = ndimage.gaussian_filter(self.image_3d, s2)
+            g1 = ndimage.gaussian_filter(image_3d, s1)
+            g2 = ndimage.gaussian_filter(image_3d, s2)
             dog = g1 - g2
             features.append(dog)
         
-        print("Computing sobel and gradients")
+        #print("Computing sobel and gradients")
         # Gradient computations using scipy
-        gx = ndimage.sobel(self.image_3d, axis=2, mode='reflect')  # x direction
-        gy = ndimage.sobel(self.image_3d, axis=1, mode='reflect')  # y direction
-        gz = ndimage.sobel(self.image_3d, axis=0, mode='reflect')  # z direction
+        gx = ndimage.sobel(image_3d, axis=2, mode='reflect')  # x direction
+        gy = ndimage.sobel(image_3d, axis=1, mode='reflect')  # y direction
+        gz = ndimage.sobel(image_3d, axis=0, mode='reflect')  # z direction
         
         # Gradient magnitude
-        print("Computing gradient magnitude")
+        #print("Computing gradient magnitude")
         gradient_magnitude = np.sqrt(gx**2 + gy**2 + gz**2)
         features.append(gradient_magnitude)
         
         # Verify shapes
-        print("Verifying shapes")
+        #print("Verifying 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
+        
+        return np.stack(features, axis=-1)
+
+    def compute_feature_maps_cpu_parallel(self, image_3d=None):
+        """Use ThreadPoolExecutor 
+        
+        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.
+        """
+        if image_3d is None:
+            image_3d = self.image_3d
+        original_shape = image_3d.shape
+        
+        features = []
+        
+        # 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 [0.5, 1.0, 2.0, 4.0]:
+                future = executor.submit(ndimage.gaussian_filter, image_3d, sigma)
+                futures.append(future)
+            
+            # Difference of Gaussians
+            for (s1, s2) in [(1, 2), (2, 4)]:
+                # Need to define a local function for this task
+                def compute_dog_local(img, s1, s2):
+                    g1 = ndimage.gaussian_filter(img, s1)
+                    g2 = ndimage.gaussian_filter(img, s2)
+                    return g1 - g2
+                
+                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)
+            
+            # Collect results
+            for future in futures:
+                result = future.result()
+                features.append(result)
+        
+        # Verify shapes
         for i, feat in enumerate(features):
             if feat.shape != original_shape:
                 feat_adjusted = np.expand_dims(feat, axis=0)
@@ -723,17 +1187,66 @@ class InteractiveSegmenter:
         return dict(z_dict)  # Convert back to regular dict
 
     def process_chunk(self, chunk_coords):
-        """Process a chunk of coordinates"""
-
+        """
+        Process a chunk of coordinates, handling both mem_lock and non-mem_lock cases.
+        Uses a consistent approach based on coordinates.
+        
+        Parameters:
+        -----------
+        chunk_coords : list of tuples
+            List of (z,y,x) coordinate tuples to process
+        
+        Returns:
+        --------
+        tuple : (foreground, background)
+            Sets of coordinates classified as foreground or background
+        """
         foreground = set()
         background = set()
-
+        
         if not self.use_two:
-
-
-            features = [self.feature_cache[z, y, x] for z, y, x in chunk_coords]
+            if self.mem_lock:
+                # For mem_lock, we need to extract a subarray and compute features
+                
+                # 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)
+                    
+            else:
+                # For non-mem_lock, simply use the feature cache
+                features = [self.feature_cache[z, y, x] for z, y, x in chunk_coords]
+            
+            # Make predictions
             predictions = self.model.predict(features)
             
+            # Assign coordinates based on predictions
             for coord, pred in zip(chunk_coords, predictions):
                 if pred:
                     foreground.add(coord)
@@ -776,45 +1289,107 @@ class InteractiveSegmenter:
         """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)
 
+        if self.mem_lock:
+            chunk_size = 32 #memory efficient 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.
+            Each chunk is a complete z-slice with all y,x coordinates,
+            unless the slice exceeds 32768 pixels, in which case it's divided into subchunks.
             
             Returns:
-                List of chunks, where each chunk contains the coordinates for one z-slice
+                List of chunks, where each chunk contains the coordinates for one z-slice or subchunk
             """
-            # Pre-calculate the number of coordinates per z-slice
-            coords_per_slice = self.image_3d.shape[1] * self.image_3d.shape[2]
-
-            # Create all coordinates at once
+            MAX_CHUNK_SIZE = 32768000
             chunks = []
+            
             for z in range(self.image_3d.shape[0]):
-                # Create y, x meshgrid once
-                y_coords, x_coords = np.meshgrid(
-                    np.arange(self.image_3d.shape[1]),
-                    np.arange(self.image_3d.shape[2]),
-                    indexing='ij'
-                )
-                
-                # Create the slice coordinates more efficiently
-                slice_coords = np.column_stack((
-                    np.full(coords_per_slice, z),
-                    y_coords.ravel(),
-                    x_coords.ravel()
-                ))
+                # 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
                 
-                # Convert to list of tuples
-                chunks.append(list(map(tuple, slice_coords)))
-
+                # If the slice is small enough, do not subchunk
+                if total_pixels <= MAX_CHUNK_SIZE or not self.mem_lock:
+                    y_coords, x_coords = np.meshgrid(
+                        np.arange(y_dim),
+                        np.arange(x_dim),
+                        indexing='ij'
+                    )
+                    
+                    slice_coords = np.column_stack((
+                        np.full(total_pixels, z),
+                        y_coords.ravel(),
+                        x_coords.ravel()
+                    ))
+                    
+                    chunks.append(list(map(tuple, slice_coords)))
+                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)
+                            y_subrange = np.arange(i, end_i)
+                            
+                            # Create meshgrid for this subchunk
+                            y_sub, x_sub = np.meshgrid(
+                                y_subrange,
+                                np.arange(x_dim),
+                                indexing='ij'
+                            )
+                            
+                            # Create coordinates for this subchunk
+                            subchunk_size = len(y_subrange) * x_dim
+                            subchunk_coords = np.column_stack((
+                                np.full(subchunk_size, z),
+                                y_sub.ravel(),
+                                x_sub.ravel()
+                            ))
+                            
+                            chunks.append(list(map(tuple, subchunk_coords)))
+                    else:  # largest_dim == 'x'
+                        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)
+                            x_subrange = np.arange(i, end_i)
+                            
+                            # Create meshgrid for this subchunk
+                            y_sub, x_sub = np.meshgrid(
+                                np.arange(y_dim),
+                                x_subrange,
+                                indexing='ij'
+                            )
+                            
+                            # Create coordinates for this subchunk
+                            subchunk_size = y_dim * len(x_subrange)
+                            subchunk_coords = np.column_stack((
+                                np.full(subchunk_size, z),
+                                y_sub.ravel(),
+                                x_sub.ravel()
+                            ))
+                            
+                            chunks.append(list(map(tuple, subchunk_coords)))
+            
             return chunks
+
         try:
             from cuml.ensemble import RandomForestClassifier as cuRandomForestClassifier
         except:
             print("Cannot find cuML, using CPU to segment instead...")
             gpu = False
         
-        if self.feature_cache is None:
+        if self.feature_cache is None and not self.mem_lock and not self.use_two:
             with self.lock:
                 if self.feature_cache is None:
                     self.feature_cache = self.compute_feature_maps()
@@ -877,23 +1452,29 @@ class InteractiveSegmenter:
 
         print("Segmenting chunks...")
 
-        
-        with ThreadPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
-            if gpu:
-                try:
-                    futures = [executor.submit(self.process_chunk_GPU, chunk) for chunk in chunks]
-                except:
-                    futures = [executor.submit(self.process_chunk, chunk) for chunk in chunks]
+        if not self.mem_lock:
+            with ThreadPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
+                if gpu:
+                    try:
+                        futures = [executor.submit(self.process_chunk_GPU, chunk) for chunk in chunks]
+                    except:
+                        futures = [executor.submit(self.process_chunk, chunk) for chunk in chunks]
 
-            else:
-                futures = [executor.submit(self.process_chunk, chunk) for chunk in chunks]
-            
-            for i, future in enumerate(futures):
-                fore, back = future.result()
+                else:
+                    futures = [executor.submit(self.process_chunk, chunk) for chunk in chunks]
+                
+                for i, future in enumerate(futures):
+                    fore, back = future.result()
+                    foreground_coords.update(fore)
+                    background_coords.update(back)
+                    print(f"Processed {i}/{len(chunks)} chunks")
+        else: #Prioritize RAM
+            for i, chunk in enumerate(chunks):
+                fore, back = self.process_chunk(chunk)
                 foreground_coords.update(fore)
                 background_coords.update(back)
-                if i % 10 == 0:
-                    print(f"Processed {i}/{len(chunks)} chunks")
+                chunk[i] = None #Help garbage collection
+                print(f"Processed {i}/{len(chunks)} chunks")
         
         return foreground_coords, background_coords
 
@@ -933,6 +1514,7 @@ class InteractiveSegmenter:
         # Update previous z
         self.prev_z = z
 
+
     def get_realtime_chunks(self, chunk_size = 64):
         print("Computing some overhead...")
 
@@ -1105,11 +1687,17 @@ class InteractiveSegmenter:
             except:
                 pass
 
-    def train_batch(self, foreground_array, speed = True, use_gpu = False, use_two = False):
+    def train_batch(self, foreground_array, speed = True, use_gpu = False, use_two = False, mem_lock = False):
         """Train directly on foreground and background arrays"""
 
         self.speed = speed
         self.cur_gpu = use_gpu
+        if mem_lock != self.mem_lock:
+            self.realtimechunks = None #dump ram
+            self.feature_cache = None
+
+
+        self.mem_lock = mem_lock
 
         if self.current_speed != speed:
             self.feature_cache = None
@@ -1197,6 +1785,83 @@ class InteractiveSegmenter:
                         background_features.append(feature_vector)
 
 
+        elif mem_lock: #Forces ram efficiency
+
+            box_size = 32
+
+            # 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
+            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)
+
         else:
 
             self.two_slices = []
@@ -1239,7 +1904,11 @@ class InteractiveSegmenter:
         y = np.hstack([np.ones(len(z_fore)), np.zeros(len(z_back))])
         
         # Train the model
-        self.model.fit(X, y)
+        try:
+            self.model.fit(X, y)
+        except:
+            print(X)
+            print(y)
 
         self.current_speed = speed
                 
@@ -1253,13 +1922,19 @@ class InteractiveSegmenter:
         if self._currently_segmenting is not None:
             return
 
-        with self.lock:
-            if speed:
+        #with self.lock <- cant remember why this was here
+        if speed:
 
+            if self.mem_lock:
+                output = self.compute_feature_maps_cpu_2d_parallel(z = z)
+            else:
                 output = self.compute_feature_maps_cpu_2d(z = z)
 
-            elif not speed:
+        elif not speed:
 
+            if self.mem_lock:
+                output = self.compute_deep_feature_maps_cpu_2d_parallel(z = z)
+            else:
                 output = self.compute_deep_feature_maps_cpu_2d(z = z)
 
         return output