nettracer3d 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

nettracer3d/segmenter.py

@@ -57,6 +57,7 @@ class InteractiveSegmenter:
 
         self.feature_cache = None
         self.lock = threading.Lock()
+        self._currently_segmenting = None
 
         # Current position attributes
         self.current_z = None
@@ -66,6 +67,119 @@ class InteractiveSegmenter:
         self.realtimechunks = None
         self.current_speed = False
 
+        # Tracking if we're using 2d or 3d segs
+        self.use_two = False
+        self.two_slices = []
+        self.speed = True
+        self.cur_gpu = False
+        self.map_slice = None
+        self.prev_z = None
+        self.previewing = False
+
+        #  flags to track state
+        self._currently_processing = False
+        self._skip_next_update = False
+        self._last_processed_slice = None
+
+    def segment_slice_chunked(self, slice_z, block_size=64):
+        """
+        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
+        
+        # Set processing flag with the slice we're processing
+        self._currently_processing = slice_z
+        
+        try:
+            
+            # First attempt to get the feature map
+            feature_map = None
+            
+            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
+            
+            # 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:
+                        predictions = self.model.predict(features)
+                        
+                        # 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
+
     def compute_deep_feature_maps_cpu(self):
         """Compute feature maps using CPU"""
         features = []
@@ -125,6 +239,67 @@ class InteractiveSegmenter:
         
         return np.stack(features, axis=-1)
 
+    def compute_deep_feature_maps_cpu_2d(self, z = None):
+        """Compute 2D feature maps using CPU"""
+        features = []
+
+        image_2d = self.image_3d[z, :, :]
+        original_shape = image_2d.shape
+        
+        # Gaussian using scipy
+        for sigma in [0.5, 1.0, 2.0, 4.0]:
+            smooth = ndimage.gaussian_filter(image_2d, sigma)
+            features.append(smooth)
+        
+        # Local statistics using scipy's convolve - adjusted for 2D
+        window_size = 5
+        kernel = np.ones((window_size, window_size)) / (window_size**2)
+        
+        # Local mean
+        local_mean = ndimage.convolve(image_2d, kernel, mode='reflect')
+        features.append(local_mean)
+        
+        # Local variance
+        mean = np.mean(image_2d)
+        local_var = ndimage.convolve((image_2d - mean)**2, kernel, mode='reflect')
+        features.append(local_var)
+        
+        # Gradient computations using scipy - adjusted axes for 2D
+        gx = ndimage.sobel(image_2d, axis=1, mode='reflect')  # x direction
+        gy = ndimage.sobel(image_2d, axis=0, mode='reflect')  # y direction
+        
+        # Gradient magnitude (2D version)
+        gradient_magnitude = np.sqrt(gx**2 + gy**2)
+        features.append(gradient_magnitude)
+        
+        # Second-order gradients
+        gxx = ndimage.sobel(gx, axis=1, mode='reflect')
+        gyy = ndimage.sobel(gy, axis=0, mode='reflect')
+        
+        # Laplacian (sum of second derivatives) - 2D version
+        laplacian = gxx + gyy
+        features.append(laplacian)
+        
+        # Hessian determinant - 2D version
+        hessian_det = gxx * gyy - ndimage.sobel(gx, axis=0, mode='reflect') * ndimage.sobel(gy, axis=1, mode='reflect')
+        features.append(hessian_det)
+        
+        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:
@@ -132,7 +307,10 @@ class InteractiveSegmenter:
         
         features = []
         image = self.image_gpu
+        image_3d = self.image_3d
         original_shape = self.image_3d.shape
+
+
         
         # Gaussian smoothing at different scales
         print("Obtaining gaussians")
@@ -150,7 +328,7 @@ class InteractiveSegmenter:
             features.append(dog)
         
         # Convert image to PyTorch tensor for gradient operations
-        image_torch = torch.from_numpy(self.image_3d).cuda()
+        image_torch = torch.from_numpy(image_3d).cuda()
         image_torch = image_torch.float().unsqueeze(0).unsqueeze(0)
         
         # Calculate required padding
@@ -159,7 +337,7 @@ class InteractiveSegmenter:
         
         # Create a single padded version with same padding
         pad = torch.nn.functional.pad(image_torch, (padding, padding, padding, padding, padding, padding), mode='replicate')
-        
+
         print("Computing sobel kernels")
 
         # Create sobel kernels
@@ -181,6 +359,8 @@ class InteractiveSegmenter:
         gradient_feature = gradient_magnitude.cpu().numpy().squeeze()
         
         features.append(gradient_feature)
+
+        print(features.shape)
         
         # Verify shapes
         for i, feat in enumerate(features):
@@ -194,6 +374,122 @@ class InteractiveSegmenter:
         
         return np.stack(features, axis=-1)
 
+    def compute_feature_maps_2d(self, z=None):
+        """Compute all feature maps for 2D images using GPU acceleration"""
+        
+        features = []
+
+        image = self.image_gpu[z, :, :]
+        image_2d = self.image_3d[z, :, :]
+        original_shape = image_2d.shape
+
+        # Gaussian smoothing at different scales
+        print("Obtaining gaussians")
+        for sigma in [0.5, 1.0, 2.0, 4.0]:
+            smooth = cp.asnumpy(self.gaussian_filter_gpu(image, sigma))
+            features.append(smooth)
+        
+        print("Obtaining diff of gaussians")
+        # Difference of Gaussians
+        for (s1, s2) in [(1, 2), (2, 4)]:
+            g1 = self.gaussian_filter_gpu(image, s1)
+            g2 = self.gaussian_filter_gpu(image, s2)
+            dog = cp.asnumpy(g1 - g2)
+            features.append(dog)
+        
+        # Convert image to PyTorch tensor for gradient operations
+        image_torch = torch.from_numpy(image_2d).cuda()
+        image_torch = image_torch.float().unsqueeze(0).unsqueeze(0)
+        
+        # Calculate required padding
+        kernel_size = 3
+        padding = kernel_size // 2
+        
+        # Create a single padded version with same padding
+        pad = torch.nn.functional.pad(image_torch, (padding, padding, padding, padding), mode='replicate')
+        
+        print("Computing sobel kernels")
+        # Create 2D sobel kernels
+        sobel_x = torch.tensor([-1, 0, 1], device='cuda').float().view(1, 1, 1, 3)
+        sobel_y = torch.tensor([-1, 0, 1], device='cuda').float().view(1, 1, 3, 1)
+        
+        # Compute gradients
+        print("Computing gradients")
+        gx = torch.nn.functional.conv2d(pad, sobel_x, padding=0)[:, :, :original_shape[0], :original_shape[1]]
+        gy = torch.nn.functional.conv2d(pad, sobel_y, padding=0)[:, :, :original_shape[0], :original_shape[1]]
+        
+        # Compute gradient magnitude (no z component in 2D)
+        print("Computing gradient mags")
+        gradient_magnitude = torch.sqrt(gx**2 + gy**2)
+        gradient_feature = gradient_magnitude.cpu().numpy().squeeze()
+        
+        features.append(gradient_feature)
+        
+        # Verify shapes
+        for i, feat in enumerate(features):
+            if feat.shape != original_shape:
+                # Create a copy of the feature to modify
+                feat_adjusted = feat
+                # Check dimensionality and expand if needed
+                if len(feat.shape) < len(original_shape):
+                    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}")
+                
+                # Update the original features list with the adjusted version
+                features[i] = feat_adjusted
+        
+        return np.stack(features, axis=-1)
+
+    def compute_feature_maps_cpu_2d(self, z = None):
+        """Compute feature maps for 2D images using CPU"""
+
+
+        features = []
+
+        image_2d = self.image_3d[z, :, :]
+        original_shape = image_2d.shape
+        
+        # Gaussian smoothing at different scales
+        for sigma in [0.5, 1.0, 2.0, 4.0]:
+            smooth = ndimage.gaussian_filter(image_2d, sigma)
+            features.append(smooth)
+        
+        # Difference of Gaussians
+        for (s1, s2) in [(1, 2), (2, 4)]:
+            g1 = ndimage.gaussian_filter(image_2d, s1)
+            g2 = ndimage.gaussian_filter(image_2d, s2)
+            dog = g1 - g2
+            features.append(dog)
+        
+        # Gradient computations using scipy - note axis changes for 2D
+        gx = ndimage.sobel(image_2d, axis=1, mode='reflect')  # x direction
+        gy = ndimage.sobel(image_2d, axis=0, mode='reflect')  # y direction
+        
+        # Gradient magnitude (no z component in 2D)
+        gradient_magnitude = np.sqrt(gx**2 + gy**2)
+        features.append(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):
         """Compute feature maps using CPU"""
         features = []
@@ -386,25 +682,112 @@ class InteractiveSegmenter:
         
         return foreground, background
 
+    def organize_by_z(self, coordinates):
+        """
+        Organizes a list of [z, y, x] coordinates into a dictionary of [y, x] coordinates grouped by z-value.
+        
+        Args:
+            coordinates: List of [z, y, x] coordinate lists
+            
+        Returns:
+            Dictionary with z-values as keys and lists of corresponding [y, x] coordinates as values
+        """
+        z_dict = {}
+        
+        for coord in coordinates:
+            z, y, x = coord  # Unpack the coordinates
+            
+            # Add the y, x coordinate to the appropriate z-value group
+            if z not in z_dict:
+                z_dict[z] = []
+            
+            z_dict[z].append((y, x))  # Store as tuple, not list, so it's hashable
+        
+        return z_dict
+
     def process_chunk(self, chunk_coords):
         """Process a chunk of coordinates"""
-        features = [self.feature_cache[z, y, x] for z, y, x in chunk_coords]
-        predictions = self.model.predict(features)
-        
+
         foreground = set()
         background = set()
-        for coord, pred in zip(chunk_coords, predictions):
-            if pred:
-                foreground.add(coord)
-            else:
-                background.add(coord)
+
+        if not self.use_two:
+
+
+            features = [self.feature_cache[z, y, x] for z, y, x in chunk_coords]
+            predictions = self.model.predict(features)
+            
+            for coord, pred in zip(chunk_coords, predictions):
+                if pred:
+                    foreground.add(coord)
+                else:
+                    background.add(coord)
+
+        else:
+            chunk_by_z = self.organize_by_z(chunk_coords)
+            for z, coords in chunk_by_z.items():
+
+                if 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 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)
         
         return foreground, background
 
+
+
     def segment_volume(self, 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)
 
+        def create_2d_chunks():
+            """
+            Create chunks by z-slices for 2D processing.
+            Each chunk is a complete z-slice with all y,x coordinates.
+            
+            Returns:
+                List of chunks, where each chunk contains the coordinates for one z-slice
+            """
+            # 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
+            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()
+                ))
+                
+                # Convert to list of tuples
+                chunks.append(list(map(tuple, slice_coords)))
+
+            return chunks
         try:
             from cuml.ensemble import RandomForestClassifier as cuRandomForestClassifier
         except:
@@ -418,52 +801,55 @@ class InteractiveSegmenter:
 
         print("Chunking data...")
         
-        # Determine optimal chunk size based on number of cores if not specified
-        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
-            
-            # 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)))
+        if not self.use_two:
+            # Determine optimal chunk size based on number of cores if not specified
+            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
+                
+                # 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
             
-            # Round to nearest multiple of 32 for better memory alignment
-            chunk_size = ((chunk_size + 15) // 32) * 32
-        
-        # 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,
-            np.arange(x_chunks) * chunk_size,
-            indexing='ij'
-        )).reshape(3, -1).T
-        
-        chunks = []
-        for z_start, y_start, x_start in chunk_starts:
-            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])
+            # 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 coordinates for this chunk efficiently
-            coords = np.stack(np.meshgrid(
-                np.arange(z_start, z_end),
-                np.arange(y_start, y_end),
-                np.arange(x_start, x_end),
+            # 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,
+                np.arange(x_chunks) * chunk_size,
                 indexing='ij'
             )).reshape(3, -1).T
             
-            chunks.append(list(map(tuple, coords)))
+            chunks = []
+            for z_start, y_start, x_start in chunk_starts:
+                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])
+                
+                # Create coordinates for this chunk efficiently
+                coords = np.stack(np.meshgrid(
+                    np.arange(z_start, z_end),
+                    np.arange(y_start, y_end),
+                    np.arange(x_start, x_end),
+                    indexing='ij'
+                )).reshape(3, -1).T
+                
+                chunks.append(list(map(tuple, coords)))
+        else:
+            chunks = create_2d_chunks()
         
         foreground_coords = set()
         background_coords = set()
@@ -471,7 +857,7 @@ class InteractiveSegmenter:
         print("Segmenting chunks...")
 
         
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
             if gpu:
                 try:
                     futures = [executor.submit(self.process_chunk_GPU, chunk) for chunk in chunks]
@@ -491,11 +877,39 @@ class InteractiveSegmenter:
         return foreground_coords, background_coords
 
     def update_position(self, z=None, x=None, y=None):
-        """Update current position for chunk prioritization"""
+        """Update current position for chunk prioritization with safeguards"""
+        
+        # Check if we should skip this update
+        if hasattr(self, '_skip_next_update') and self._skip_next_update:
+            self._skip_next_update = False
+            return
+        
+        # Store the previous z-position if not set
+        if not hasattr(self, 'prev_z') or self.prev_z is None:
+            self.prev_z = z
+        
+        # Check if currently processing - if so, only update position but don't trigger map_slice changes
+        if hasattr(self, '_currently_processing') and self._currently_processing:
+            self.current_z = z
+            self.current_x = x
+            self.current_y = y
+            self.prev_z = z
+            return
+        
+        # Update current positions
         self.current_z = z
         self.current_x = x
         self.current_y = y
-
+        
+        # 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.current_z not in self.feature_cache:
+                self.map_slice = None
+            self._currently_segmenting = None
+        
+        # Update previous z
+        self.prev_z = z
 
     def get_realtime_chunks(self, chunk_size = 64):
         print("Computing some overhead...")
@@ -572,6 +986,7 @@ class InteractiveSegmenter:
             gpu = False
 
 
+
         if self.realtimechunks is None:
             self.get_realtime_chunks()
         else:
@@ -658,55 +1073,141 @@ class InteractiveSegmenter:
                     del futures[future]
                     yield fore, back
 
+
     def cleanup(self):
         """Clean up GPU memory"""
         if self.use_gpu:
             cp.get_default_memory_pool().free_all_blocks()
             torch.cuda.empty_cache()
 
-    def train_batch(self, foreground_array, speed = True, use_gpu = False):
+    def train_batch(self, foreground_array, speed = True, use_gpu = False, use_two = False):
         """Train directly on foreground and background arrays"""
 
+        self.speed = speed
+        self.cur_gpu = use_gpu
+
         if self.current_speed != speed:
             self.feature_cache = None
 
-        if self.feature_cache is None:
-            with self.lock:
-                if self.feature_cache is None and speed:
-                    if use_gpu:
-                        self.feature_cache = self.compute_feature_maps()
-                    else:
-                        self.feature_cache = self.compute_feature_maps_cpu()
+        if use_two:
 
-                elif self.feature_cache is None and not speed:
-                    if use_gpu:
+            changed = [] #Track which slices need feature maps
 
-                        self.feature_cache = self.compute_deep_feature_maps()
-                    else:
-                        self.feature_cache = self.compute_deep_feature_maps_cpu()
+            if not self.use_two: #Clarifies if we need to redo feature cache for 2D
+                self.feature_cache = None
+                self.use_two = True
 
+            if self.feature_cache == None:
+                self.feature_cache = {}
 
-        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]
-            
-            # 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
-            self.model.fit(X, y)
 
-            self.current_speed = speed
-        except:
-            print("Features maps computed, but no segmentation examples were provided so the model was not trained")
+            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 = []
+
+            for i, z in enumerate(z_fore):
+                if z in changed:  # Means this slice needs a feature map
+                    new_map = self.get_feature_map_slice(z, speed, use_gpu)
+                    self.feature_cache[z] = new_map
+                    changed.remove(z)
+                
+                current_map = self.feature_cache[z]
+                
+                # Get the feature vector for this foreground point
+                feature_vector = current_map[y_fore[i], x_fore[i]]
+                
+                # Add to our collection
+                foreground_features.append(feature_vector)
+
+            for i, z in enumerate(z_back):
+                if z in changed:  # Means this slice needs a feature map
+                    new_map = self.get_feature_map_slice(z, speed, use_gpu)
+                    self.feature_cache[z] = new_map
+                
+                current_map = self.feature_cache[z]
+                
+                # Get the feature vector for this foreground point
+                feature_vector = current_map[y_back[i], x_back[i]]
+                
+                # Add to our collection
+                background_features.append(feature_vector)
+
+
+        else:
+
+            self.two_slices = []
+
+            if self.use_two: #Clarifies if we need to redo feature cache for 3D
+
+                self.feature_cache = None
+                self.use_two = False
+
+            if self.feature_cache is None:
+                with self.lock:
+                    if self.feature_cache is None and speed:
+                        if use_gpu:
+                            self.feature_cache = self.compute_feature_maps()
+                        else:
+                            self.feature_cache = self.compute_feature_maps_cpu()
+
+                    elif self.feature_cache is None and not speed:
+                        if use_gpu:
+
+                            self.feature_cache = self.compute_deep_feature_maps()
+                        else:
+                            self.feature_cache = self.compute_deep_feature_maps_cpu()
+
+
+            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:
+                print("Features maps computed, but no segmentation examples were provided so the model was not trained")
+
+
+        # 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
+        self.model.fit(X, y)
+
+        self.current_speed = speed
+                
+
 
 
         print("Done")
 
+    def get_feature_map_slice(self, z, speed, use_gpu):
+
+        if self._currently_segmenting is not None:
+            return
+
+        with self.lock:
+            if speed:
+
+                output = self.compute_feature_maps_cpu_2d(z = z)
+
+            elif not speed:
+
+                output = self.compute_deep_feature_maps_cpu_2d(z = z)
+
+        return output