nettracer3d 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

nettracer3d/segmenter.py

@@ -21,7 +21,7 @@ import multiprocessing
 from collections import defaultdict
 
 class InteractiveSegmenter:
-    def __init__(self, image_3d, use_gpu=True):
+    def __init__(self, image_3d, use_gpu=False):
         self.image_3d = image_3d
         self.patterns = []
 
@@ -82,7 +82,13 @@ class InteractiveSegmenter:
         self._last_processed_slice = None
         self.mem_lock = False
 
-    def segment_slice_chunked(self, slice_z, block_size=64):
+        #Adjustable feature map params:
+        self.alphas = [1,2,4,8]
+        self.windows = 10
+        self.dogs = [(1, 2), (2, 4), (4, 8)]
+        self.master_chunk = 49
+
+    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.
@@ -280,13 +286,24 @@ class InteractiveSegmenter:
             def compute_gaussian(sigma):
                 return ndimage.gaussian_filter(image_3d, sigma)
             
-            for sigma in [0.5, 1.0, 2.0, 4.0]:
+            for sigma in self.alphas:
                 future = executor.submit(compute_gaussian, sigma)
                 futures.append(('gaussian', sigma, future))
+
+            def compute_dog_local(img, s1, s2):
+                g1 = ndimage.gaussian_filter(img, s1)
+                g2 = ndimage.gaussian_filter(img, s2)
+                return g1 - g2
+
+            # 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 = 5
+                window_size = self.windows
                 kernel = np.ones((window_size, window_size, window_size)) / (window_size**3)
                 return ndimage.convolve(image_3d, kernel, mode='reflect')
             
@@ -294,7 +311,7 @@ class InteractiveSegmenter:
             futures.append(('local_mean', None, future))
             
             def compute_local_variance():
-                window_size = 5
+                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')
@@ -395,8 +412,11 @@ class InteractiveSegmenter:
         features = []
         
         # Add Gaussian features
-        for sigma in [0.5, 1.0, 2.0, 4.0]:
+        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'])
@@ -494,13 +514,24 @@ class InteractiveSegmenter:
             def compute_gaussian(sigma):
                 return ndimage.gaussian_filter(image_2d, sigma)
             
-            for sigma in [0.5, 1.0, 2.0, 4.0]:
+            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 = 5
+                window_size = self.windows
                 kernel = np.ones((window_size, window_size)) / (window_size**2)
                 return ndimage.convolve(image_2d, kernel, mode='reflect')
             
@@ -508,7 +539,7 @@ class InteractiveSegmenter:
             futures.append(('local_mean', None, future))
             
             def compute_local_variance():
-                window_size = 5
+                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')
@@ -608,8 +639,11 @@ class InteractiveSegmenter:
         features = []
         
         # Add Gaussian features
-        for sigma in [0.5, 1.0, 2.0, 4.0]:
+        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'])
@@ -843,7 +877,7 @@ class InteractiveSegmenter:
             def compute_gaussian(sigma):
                 return ndimage.gaussian_filter(image_2d, sigma)
             
-            gaussian_sigmas = [0.5, 1.0, 2.0, 4.0]
+            gaussian_sigmas = self.alphas
             for sigma in gaussian_sigmas:
                 future = executor.submit(compute_gaussian, sigma)
                 futures.append(('gaussian', sigma, future))
@@ -854,7 +888,7 @@ class InteractiveSegmenter:
                 g2 = ndimage.gaussian_filter(image_2d, s2)
                 return g1 - g2
             
-            dog_pairs = [(1, 2), (2, 4)]
+            dog_pairs = self.dogs
             for (s1, s2) in dog_pairs:
                 future = executor.submit(compute_dog, s1, s2)
                 futures.append(('dog', (s1, s2), future))
@@ -978,17 +1012,17 @@ class InteractiveSegmenter:
             futures = []
             
             # Gaussian smoothing at different scales
-            for sigma in [0.5, 1.0, 2.0, 4.0]:
+            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)
+                g2 = ndimage.gaussian_filter(img, s2)
+                return g1 - g2
+
             # 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
+            for (s1, s2) in self.dogs:
                 
                 future = executor.submit(compute_dog_local, image_3d, s1, s2)
                 futures.append(future)
@@ -1180,9 +1214,10 @@ class InteractiveSegmenter:
             Dictionary with z-values as keys and lists of corresponding [y, x] coordinates as values
         """
         z_dict = defaultdict(list)
-        
+
         for z, y, x in coordinates:
             z_dict[z].append((y, x))
+
         
         return dict(z_dict)  # Convert back to regular dict
 
@@ -1204,19 +1239,34 @@ class InteractiveSegmenter:
         foreground = set()
         background = set()
         
-        if not self.use_two:
+        if self.previewing or not self.use_two:
             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)
-                
+
+                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]
@@ -1254,8 +1304,13 @@ class InteractiveSegmenter:
                     background.add(coord)
 
         else:
-            chunk_by_z = self.organize_by_z(chunk_coords)
-            for z, coords in chunk_by_z.items():
+
+            if self.mem_lock:
+                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)
@@ -1283,6 +1338,63 @@ class InteractiveSegmenter:
         
         return foreground, background
 
+    def twodim_coords(self, y_dim, x_dim, z, chunk_size = None, subrange = None):
+
+        if subrange is None:
+            y_coords, x_coords = np.meshgrid(
+                np.arange(y_dim),
+                np.arange(x_dim),
+                indexing='ij'
+            )
+        
+            slice_coords = np.column_stack((
+                np.full(chunk_size, z),
+                y_coords.ravel(),
+                x_coords.ravel()
+            ))
+
+        elif subrange[0] == 'y':
+
+            y_subrange = np.arange(subrange[1], subrange[2])
+
+            # 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
+            slice_coords = np.column_stack((
+                np.full(subchunk_size, z),
+                y_sub.ravel(),
+                x_sub.ravel()
+            ))
+
+        elif subrange[0] == 'x':
+
+            x_subrange = np.arange(subrange[1], subrange[2])
+            
+            # 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)
+            slice_coords = np.column_stack((
+                np.full(subchunk_size, z),
+                y_sub.ravel(),
+                x_sub.ravel()
+            ))
+
+
+
+        return list(map(tuple, slice_coords))
+        
 
 
     def segment_volume(self, chunk_size=None, gpu=False):
@@ -1293,19 +1405,21 @@ class InteractiveSegmenter:
         self.map_slice = None
 
         if self.mem_lock:
-            chunk_size = 32 #memory efficient chunk
+            chunk_size = self.master_chunk #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,
-            unless the slice exceeds 32768 pixels, in which case it's divided into subchunks.
+            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 = 32768
+            MAX_CHUNK_SIZE = 262144
+            if not self.mem_lock:
+                MAX_CHUNK_SIZE = 10000000000000000000000000 #unlimited i guess
             chunks = []
             
             for z in range(self.image_3d.shape[0]):
@@ -1315,20 +1429,16 @@ class InteractiveSegmenter:
                 total_pixels = y_dim * x_dim
                 
                 # 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)))
+                if total_pixels <= MAX_CHUNK_SIZE:
+
+
+                    if not self.mem_lock:
+                        chunks.append(self.twodim_coords(y_dim, x_dim, z, total_pixels))
+                    else:
+                        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'
@@ -1342,56 +1452,31 @@ class InteractiveSegmenter:
                         # 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)))
+
+                            if not self.mem_lock:
+                                chunks.append(self.twodim_coords(y_dim, x_dim, z, None, ['y', i, end_i]))
+                            else:
+                                chunks.append([y_dim, x_dim, z, None, ['y', i, end_i]])
+
                     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)))
+
+                            if not self.mem_lock:
+                                chunks.append(self.twodim_coords(y_dim, x_dim, z, None, ['x', i, end_i]))
+                            else:
+                                chunks.append([y_dim, x_dim, z, None, ['x', i, end_i]])
             
             return chunks
 
-        try:
-            from cuml.ensemble import RandomForestClassifier as cuRandomForestClassifier
-        except:
-            print("Cannot find cuML, using CPU to segment instead...")
-            gpu = False
-        
+        #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 and not self.mem_lock and not self.use_two:
             with self.lock:
                 if self.feature_cache is None:
@@ -1437,15 +1522,24 @@ class InteractiveSegmenter:
                 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
+                if self.mem_lock:
+                    # Create coordinates for this chunk efficiently
+                    coords = [z_start, z_end, y_start, y_end, x_start, x_end]
+                    chunks.append(coords)
+
+                else:
+                    # Consider moving this to process chunk ??
+                    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)))
+                    chunks.append(list(map(tuple, coords)))
+
+
+
         else:
             chunks = create_2d_chunks()
             self.feature_cache = None #Decided this should not maintain training data for segmenting 2D
@@ -1476,7 +1570,10 @@ class InteractiveSegmenter:
                 fore, back = self.process_chunk(chunk)
                 foreground_coords.update(fore)
                 background_coords.update(back)
-                chunk[i] = None #Help garbage collection
+                try:
+                    chunk[i] = None #Help garbage collection
+                except:
+                    pass
                 print(f"Processed {i}/{len(chunks)} chunks")
         
         return foreground_coords, background_coords
@@ -1518,10 +1615,7 @@ class InteractiveSegmenter:
         self.prev_z = z
 
 
-    def get_realtime_chunks(self, chunk_size = 64):
-        print("Computing some overhead...")
-
-
+    def get_realtime_chunks(self, chunk_size = 49):
 
         # Determine if we need to chunk XY planes
         small_dims = (self.image_3d.shape[1] <= chunk_size and 
@@ -1544,16 +1638,9 @@ class InteractiveSegmenter:
         # Create chunks for each Z plane
         for z in range(self.image_3d.shape[0]):
             if small_dims:
-                # One chunk per Z
-                coords = np.stack(np.meshgrid(
-                    [z],
-                    np.arange(self.image_3d.shape[1]),
-                    np.arange(self.image_3d.shape[2]),
-                    indexing='ij'
-                )).reshape(3, -1).T
                 
                 chunk_dict[(z, 0, 0)] = {
-                    'coords': list(map(tuple, coords)),
+                    'coords': [0, self.image_3d.shape[1], 0, self.image_3d.shape[2]],
                     'processed': False,
                     'z': z
                 }
@@ -1566,15 +1653,8 @@ class InteractiveSegmenter:
                         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])
                         
-                        coords = np.stack(np.meshgrid(
-                            [z],
-                            np.arange(y_start, y_end),
-                            np.arange(x_start, x_end),
-                            indexing='ij'
-                        )).reshape(3, -1).T
-                        
                         chunk_dict[(z, y_start, x_start)] = {
-                            'coords': list(map(tuple, coords)),
+                            'coords': [y_start, y_end, x_start, x_end],
                             'processed': False,
                             'z': z
                         }
@@ -1606,15 +1686,15 @@ class InteractiveSegmenter:
         def get_nearest_unprocessed_chunk(self):
             """Get nearest unprocessed chunk prioritizing current Z"""
             curr_z = self.current_z if self.current_z is not None else self.image_3d.shape[0] // 2
-            curr_y = self.current_x if self.current_x is not None else self.image_3d.shape[1] // 2
-            curr_x = self.current_y if self.current_y is not None else self.image_3d.shape[2] // 2
+            curr_y = self.current_y if self.current_y is not None else self.image_3d.shape[1] // 2
+            curr_x = self.current_x if self.current_x is not None else self.image_3d.shape[2] // 2
             
             # First try to find chunks at current Z
             current_z_chunks = [(pos, info) for pos, info in chunk_dict.items() 
-                              if info['z'] == curr_z and not info['processed']]
+                              if pos[0] == curr_z and not info['processed']]
             
             if current_z_chunks:
-                # Find nearest chunk in current Z plane
+                # Find nearest chunk in current Z plane using the chunk positions from the key
                 nearest = min(current_z_chunks, 
                             key=lambda x: ((x[0][1] - curr_y) ** 2 + 
                                          (x[0][2] - curr_x) ** 2))
@@ -1631,7 +1711,7 @@ class InteractiveSegmenter:
                 target_z = available_z[0][0]
                 # Find nearest chunk in target Z plane
                 z_chunks = [(pos, info) for pos, info in chunk_dict.items() 
-                           if info['z'] == target_z and not info['processed']]
+                           if pos[0] == target_z and not info['processed']]
                 nearest = min(z_chunks, 
                             key=lambda x: ((x[0][1] - curr_y) ** 2 + 
                                          (x[0][2] - curr_x) ** 2))
@@ -1640,45 +1720,38 @@ class InteractiveSegmenter:
             return None
         
 
-        with ThreadPoolExecutor() as executor:
-            futures = {}
-            import multiprocessing
-            total_cores = multiprocessing.cpu_count()
-            #available_workers = max(1, min(4, total_cores // 2))  # Use half cores, max of 4
-            available_workers = 1
+        while True:
+            # Find nearest unprocessed chunk using class attributes
+            chunk_idx = get_nearest_unprocessed_chunk(self)
+            if chunk_idx is None:
+                break
+                
+            # Process the chunk directly
+            chunk = chunk_dict[chunk_idx]
+            chunk['processed'] = True
+            coords = chunk['coords']
+
+            coords = np.stack(np.meshgrid(
+                [chunk['z']],
+                np.arange(coords[0], coords[1]),
+                np.arange(coords[2], coords[3]),
+                indexing='ij'
+            )).reshape(3, -1).T
 
-            while True:
-                # Find nearest unprocessed chunk using class attributes
-                chunk_idx = get_nearest_unprocessed_chunk(self)  # Pass self
-                if chunk_idx is None:
-                    break
-                    
-                while (len(futures) < available_workers and 
-                       (chunk_idx := get_nearest_unprocessed_chunk(self))):  # Pass self
-                    chunk = chunk_dict[chunk_idx]
-                    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:
-                        future = executor.submit(self.process_chunk, chunk['coords'])
+            coords = list(map(tuple, coords))
 
-                    futures[future] = chunk_idx
-                    chunk['processed'] = True
-                
-                # Check completed futures
-                done, _ = concurrent.futures.wait(
-                    futures.keys(),
-                    timeout=0.1,
-                    return_when=concurrent.futures.FIRST_COMPLETED
-                )
-                
-                # Process completed chunks
-                for future in done:
-                    fore, back = future.result()
-                    del futures[future]
-                    yield fore, back
+            
+            # Process the chunk directly based on whether GPU is available
+            if gpu:
+                try:
+                    fore, back = self.process_chunk_GPU(coords)
+                except:
+                    fore, back = self.process_chunk(coords)
+            else:
+                fore, back = self.process_chunk(coords)
+            
+            # Yield the results
+            yield fore, back
 
 
     def cleanup(self):
@@ -1700,6 +1773,8 @@ class InteractiveSegmenter:
             self.realtimechunks = None #dump ram
             self.feature_cache = None
 
+        if not use_two:
+            self.use_two = False
 
         self.mem_lock = mem_lock
 
@@ -1791,7 +1866,7 @@ class InteractiveSegmenter:
 
         elif mem_lock: #Forces ram efficiency
 
-            box_size = 32
+            box_size = self.master_chunk
 
             # Memory-efficient approach: compute features only for necessary subarrays
             foreground_features = []