nettracer3d 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

nettracer3d/nettracer_gui.py

@@ -397,6 +397,9 @@ class ImageViewerWindow(QMainWindow):
         # Initialize highlight overlay
         self.highlight_overlay = None
         self.highlight_bounds = None  # Store bounds for positioning
+        self.mini_overlay = False # If the program is currently drawing the overlay by frame this will be true
+        self.mini_overlay_data = None #Actual data for mini overlay
+        self.mini_thresh = (500*500*500) # Array volume to start using mini overlays for
 
     def start_left_scroll(self):
         """Start scrolling left when left arrow is pressed."""
@@ -443,6 +446,9 @@ class ImageViewerWindow(QMainWindow):
             edge_indices (list): List of edge indices to highlight
         """
 
+        self.mini_overlay = False #If this method is ever being called, it means we are rendering the entire overlay so mini overlay needs to reset.
+        self.mini_overlay_data = None
+
         def process_chunk(chunk_data, indices_to_check):
             """Process a single chunk of the array to create highlight mask"""
             mask = np.isin(chunk_data, indices_to_check)
@@ -463,6 +469,9 @@ class ImageViewerWindow(QMainWindow):
         if overlay1_indices is not None:
             if 0 in overlay1_indices:
                 overlay1_indices.remove(0)
+        if overlay2_indices is not None:
+            if 0 in overlay2_indices:
+                overlay2_indices.remove(0)
 
         if node_indices is None:
             node_indices = []
@@ -555,6 +564,8 @@ class ImageViewerWindow(QMainWindow):
 
     def create_highlight_overlay_slice(self, indices, bounds = False):
 
+        """Highlight overlay generation method specific for the segmenter interactive mode"""
+
 
         def process_chunk_bounds(chunk_data, indices_to_check):
             """Process a single chunk of the array to create highlight mask"""
@@ -628,6 +639,110 @@ class ImageViewerWindow(QMainWindow):
         # Update display
         self.update_display(preserve_zoom=(current_xlim, current_ylim), called = True)
 
+    def create_mini_overlay(self, node_indices = None, edge_indices = None):
+
+        """
+        Create a highlight overlay one slice at a time.
+        
+        Args:
+            node_indices (list): List of node indices to highlight
+            edge_indices (list): List of edge indices to highlight
+        """
+
+        def process_chunk(chunk_data, indices_to_check):
+            """Process a single chunk of the array to create highlight mask"""
+            mask = np.isin(chunk_data, indices_to_check)
+            return mask * 255
+
+
+        if node_indices is not None:
+            if 0 in node_indices:
+                node_indices.remove(0)
+        if edge_indices is not None:
+            if 0 in edge_indices:
+                edge_indices.remove(0)
+
+
+        if node_indices is None:
+            node_indices = []
+        if edge_indices is None:
+            edge_indices = []
+
+            
+        current_xlim = self.ax.get_xlim() if hasattr(self, 'ax') and self.ax.get_xlim() != (0, 1) else None
+        current_ylim = self.ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
+        
+        if not node_indices and not edge_indices: #Theoretically this can't be called because it uses full highlight overlay method for empty clicks
+            self.mini_overlay_data = None
+            self.mini_overlay = False
+            self.update_display(preserve_zoom=(current_xlim, current_ylim))
+            return
+            
+        # Get the shape of the mini array from any existing channel
+        for channel in self.channel_data:
+            if channel is not None:
+                full_shape = channel.shape
+                full_shape = (full_shape[1], full_shape[2]) #Just get (Y, X) shape
+                break
+        else:
+            return  # No valid channels to get shape from
+            
+        # Initialize full-size overlay
+        self.mini_overlay_data = np.zeros(full_shape, dtype=np.uint8)
+        
+        # Get number of CPU cores
+        num_cores = mp.cpu_count()
+        
+        # Calculate chunk size along y-axis
+        chunk_size = full_shape[0] // num_cores
+        if chunk_size < 1:
+            chunk_size = 1
+        
+        def process_channel(channel_data, indices, array_shape):
+            if channel_data is None or not indices:
+                return None
+                
+            # Create chunks
+            chunks = []
+            for i in range(0, array_shape[0], chunk_size):
+                end = min(i + chunk_size, array_shape[0])
+                chunks.append(channel_data[i:end, :])
+                
+            # Process chunks in parallel using ThreadPoolExecutor
+            process_func = partial(process_chunk, indices_to_check=indices)
+
+            
+            with ThreadPoolExecutor(max_workers=num_cores) as executor:
+                chunk_results = list(executor.map(process_func, chunks))
+                
+            # Reassemble the chunks
+            return np.concatenate(chunk_results, axis=0)
+        
+        # Process nodes and edges in parallel using multiprocessing
+        with ThreadPoolExecutor(max_workers=num_cores) as executor:
+            try:
+                slice_node = self.channel_data[0][self.current_slice, :, :] #This is the only major difference to the big highlight... we are only looking at this
+                future_nodes = executor.submit(process_channel, slice_node, node_indices, full_shape)
+                node_overlay = future_nodes.result()
+            except:
+                node_overlay = None
+            try:
+                slice_edge = self.channel_data[1][self.current_slice, :, :]
+                future_edges = executor.submit(process_channel, slice_edge, edge_indices, full_shape)
+                edge_overlay = future_edges.result()
+            except:
+                edge_overlay = None
+                        
+        # Combine results
+        if node_overlay is not None:
+            self.mini_overlay_data = np.maximum(self.mini_overlay_data, node_overlay)
+        if edge_overlay is not None:
+            self.mini_overlay_data = np.maximum(self.mini_overlay_data, edge_overlay)
+
+                
+        # Update display
+        self.update_display(preserve_zoom=(current_xlim, current_ylim))
+
 
 
         
@@ -890,14 +1005,23 @@ class ImageViewerWindow(QMainWindow):
                     edge_indices = filtered_df.iloc[:, 2].unique().tolist()
                     self.clicked_values['edges'] = edge_indices
 
-                    self.create_highlight_overlay(
-                        node_indices=self.clicked_values['nodes'], 
-                        edge_indices=self.clicked_values['edges']
-                    )
+                    if self.channel_data[1].shape[0] * self.channel_data[1].shape[1] * self.channel_data[1].shape[2] > self.mini_thresh:
+                        self.mini_overlay = True
+                        self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                    else:
+                        self.create_highlight_overlay(
+                            node_indices=self.clicked_values['nodes'], 
+                            edge_indices=self.clicked_values['edges']
+                        )
                 else:
-                    self.create_highlight_overlay(
-                        node_indices=self.clicked_values['nodes']
-                )
+                    if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                        self.mini_overlay = True
+                        self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                    else:
+                        self.create_highlight_overlay(
+                            node_indices=self.clicked_values['nodes'],
+                            edge_indices = self.clicked_values['edges']
+                    )
             
                 
         except Exception as e:
@@ -972,14 +1096,23 @@ class ImageViewerWindow(QMainWindow):
             if edges:
                 edge_indices = filtered_df.iloc[:, 2].unique().tolist()
                 self.clicked_values['edges'] = edge_indices
-                self.create_highlight_overlay(
-                    node_indices=self.clicked_values['nodes'],
-                    edge_indices=edge_indices
-                )
+                if self.channel_data[1].shape[0] * self.channel_data[1].shape[1] * self.channel_data[1].shape[2] > self.mini_thresh:
+                    self.mini_overlay = True
+                    self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                else:
+                    self.create_highlight_overlay(
+                        node_indices=self.clicked_values['nodes'],
+                        edge_indices=edge_indices
+                    )
             else:
-                self.create_highlight_overlay(
-                    node_indices = self.clicked_values['nodes']
-            )
+                if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                    self.mini_overlay = True
+                    self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                else:
+                    self.create_highlight_overlay(
+                        node_indices = self.clicked_values['nodes'],
+                        edge_indices = self.clicked_values['edges']
+                )
 
         except Exception as e:
 
@@ -1043,15 +1176,24 @@ class ImageViewerWindow(QMainWindow):
                 if edges:
                     edge_indices = filtered_df.iloc[:, 2].unique().tolist()
                     self.clicked_values['edges'] = edge_indices
-                    self.create_highlight_overlay(
-                        node_indices=nodes,
-                        edge_indices=edge_indices
-                    )
+                    if self.channel_data[1].shape[0] * self.channel_data[1].shape[1] * self.channel_data[1].shape[2] > self.mini_thresh:
+                        self.mini_overlay = True
+                        self.create_mini_overlay(node_indices = nodes, edge_indices = edge_indices)
+                    else:
+                        self.create_highlight_overlay(
+                            node_indices=nodes,
+                            edge_indices=edge_indices
+                        )
                     self.clicked_values['nodes'] = nodes
                 else:
-                    self.create_highlight_overlay(
-                        node_indices = nodes
-                )
+                    if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                        self.mini_overlay = True
+                        self.create_mini_overlay(node_indices = nodes, edge_indices = self.clicked_values['edges'])
+                    else:
+                        self.create_highlight_overlay(
+                            node_indices = nodes,
+                            edge_indices = self.clicked_values['edges']
+                    )
                     self.clicked_values['nodes'] = nodes
 
         except Exception as e:
@@ -1095,9 +1237,15 @@ class ImageViewerWindow(QMainWindow):
 
             print(f"Found {len(filtered_df)} direct connections between nodes of ID {sort} and their neighbors (of any ID)")
 
-            self.create_highlight_overlay(
-                node_indices= nodes
-                )
+            if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                self.mini_overlay = True
+                self.create_mini_overlay(node_indices = nodes, edge_indices = self.clicked_values['edges'])
+            else:
+                self.create_highlight_overlay(
+                    node_indices = nodes,
+                    edge_indices = self.clicked_values['edges']
+                    )
+            self.clicked_values['nodes'] = nodes
 
         except Exception as e:
             print(f"Error showing identities: {e}")
@@ -1393,8 +1541,6 @@ class ImageViewerWindow(QMainWindow):
 
             pairs = list(combinations(nodes, 2))
 
-            print(pairs)
-
 
             for i in range(len(my_network.network_lists[0]) - 1, -1, -1):
                 print((my_network.network_lists[0][i], my_network.network_lists[1][i]))
@@ -1891,7 +2037,11 @@ class ImageViewerWindow(QMainWindow):
                         self.clicked_values['nodes'].extend(selected_values)
                         # Remove duplicates while preserving order
                         self.clicked_values['nodes'] = list(dict.fromkeys(self.clicked_values['nodes']))
-                        self.create_highlight_overlay(node_indices=self.clicked_values['nodes'])
+                        if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                            self.mini_overlay = True
+                            self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                        else:
+                            self.create_highlight_overlay(node_indices=self.clicked_values['nodes'])
                         
                         # Try to highlight the last selected value in tables
                         if self.clicked_values['nodes']:
@@ -1903,7 +2053,11 @@ class ImageViewerWindow(QMainWindow):
                         self.clicked_values['edges'].extend(selected_values)
                         # Remove duplicates while preserving order
                         self.clicked_values['edges'] = list(dict.fromkeys(self.clicked_values['edges']))
-                        self.create_highlight_overlay(edge_indices=self.clicked_values['edges'])
+                        if self.channel_data[1].shape[0] * self.channel_data[1].shape[1] * self.channel_data[1].shape[2] > self.mini_thresh:
+                            self.mini_overlay = True
+                            self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                        else:
+                            self.create_highlight_overlay(edge_indices=self.clicked_values['edges'])
                         
                         # Try to highlight the last selected value in tables
                         if self.clicked_values['edges']:
@@ -2144,9 +2298,17 @@ class ImageViewerWindow(QMainWindow):
 
                     # Highlight the clicked element in the image using the stored lists                
                     if self.active_channel == 0 and (starting_vals['nodes']) != (self.clicked_values['nodes']):
-                        self.create_highlight_overlay(node_indices=self.clicked_values['nodes'], edge_indices=self.clicked_values['edges'])
+                        if self.channel_data[0].shape[0] * self.channel_data[0].shape[1] * self.channel_data[0].shape[2] > self.mini_thresh:
+                            self.mini_overlay = True
+                            self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                        else:
+                            self.create_highlight_overlay(node_indices=self.clicked_values['nodes'], edge_indices=self.clicked_values['edges'])
                     elif self.active_channel == 1 and starting_vals['edges'] != self.clicked_values['edges']:
-                        self.create_highlight_overlay(node_indices=self.clicked_values['nodes'], edge_indices=self.clicked_values['edges'])
+                        if self.channel_data[1].shape[0] * self.channel_data[1].shape[1] * self.channel_data[1].shape[2] > self.mini_thresh:
+                            self.mini_overlay = True
+                            self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
+                        else:
+                            self.create_highlight_overlay(node_indices=self.clicked_values['nodes'], edge_indices=self.clicked_values['edges'])
 
                                 
                 except IndexError:
@@ -2298,6 +2460,8 @@ class ImageViewerWindow(QMainWindow):
         netoverlay_action.triggered.connect(self.show_netoverlay_dialog)
         idoverlay_action = overlay_menu.addAction("Create ID Overlay")
         idoverlay_action.triggered.connect(self.show_idoverlay_dialog)
+        coloroverlay_action = overlay_menu.addAction("Color Nodes (or Edges)")
+        coloroverlay_action.triggered.connect(self.show_coloroverlay_dialog)
         searchoverlay_action = overlay_menu.addAction("Show Search Regions")
         searchoverlay_action.triggered.connect(self.show_search_dialog)
         shuffle_action = overlay_menu.addAction("Shuffle")
@@ -2578,6 +2742,11 @@ class ImageViewerWindow(QMainWindow):
         dialog = IdOverlayDialog(self)
         dialog.exec()
 
+    def show_coloroverlay_dialog(self):
+        """show the color overlay dialog"""
+        dialog = ColorOverlayDialog(self)
+        dialog.exec()
+
     def show_search_dialog(self):
         """Show the search dialog"""
         dialog = SearchOverlayDialog(self)
@@ -3001,14 +3170,20 @@ class ImageViewerWindow(QMainWindow):
             else:
                 self.channel_data[channel_index] = channel_data
 
-            if len(self.channel_data[channel_index].shape) == 3:  # potentially 2D RGB
-                if self.channel_data[channel_index].shape[-1] in (3, 4):  # last dim is 3 or 4
-                    if self.confirm_rgb_dialog():
-                        # User confirmed it's 2D RGB, expand to 4D
-                        self.channel_data[channel_index] = np.expand_dims(self.channel_data[channel_index], axis=0)
-                        
-            if len(self.channel_data[channel_index].shape) == 4 and (channel_index == 0 or channel_index == 1):
-                self.channel_data[channel_index] = self.reduce_rgb_dimension(self.channel_data[channel_index])
+            try:
+                if len(self.channel_data[channel_index].shape) == 3:  # potentially 2D RGB
+                    if self.channel_data[channel_index].shape[-1] in (3, 4):  # last dim is 3 or 4
+                        if self.confirm_rgb_dialog():
+                            # User confirmed it's 2D RGB, expand to 4D
+                            self.channel_data[channel_index] = np.expand_dims(self.channel_data[channel_index], axis=0)
+            except:
+                pass
+            
+            try:
+                if len(self.channel_data[channel_index].shape) == 4 and (channel_index == 0 or channel_index == 1):
+                    self.channel_data[channel_index] = self.reduce_rgb_dimension(self.channel_data[channel_index])
+            except:
+                pass
 
             reset_resize = False
 
@@ -3050,27 +3225,30 @@ class ImageViewerWindow(QMainWindow):
                 self.active_channel_combo.setEnabled(True)
             
             # Update slider range if this is the first channel loaded
-            if len(self.channel_data[channel_index].shape) == 3 or len(self.channel_data[channel_index].shape) == 4:
-                if not self.slice_slider.isEnabled():
-                    self.slice_slider.setEnabled(True)
-                    self.slice_slider.setMinimum(0)
-                    self.slice_slider.setMaximum(self.channel_data[channel_index].shape[0] - 1)
-                    if self.slice_slider.value() < self.channel_data[channel_index].shape[0] - 1:
-                        self.current_slice = self.slice_slider.value()
+            try:
+                if len(self.channel_data[channel_index].shape) == 3 or len(self.channel_data[channel_index].shape) == 4:
+                    if not self.slice_slider.isEnabled():
+                        self.slice_slider.setEnabled(True)
+                        self.slice_slider.setMinimum(0)
+                        self.slice_slider.setMaximum(self.channel_data[channel_index].shape[0] - 1)
+                        if self.slice_slider.value() < self.channel_data[channel_index].shape[0] - 1:
+                            self.current_slice = self.slice_slider.value()
+                        else:
+                            self.slice_slider.setValue(0)
+                            self.current_slice = 0
                     else:
-                        self.slice_slider.setValue(0)
-                        self.current_slice = 0
+                        self.slice_slider.setEnabled(True)
+                        self.slice_slider.setMinimum(0)
+                        self.slice_slider.setMaximum(self.channel_data[channel_index].shape[0] - 1)
+                        if self.slice_slider.value() < self.channel_data[channel_index].shape[0] - 1:
+                            self.current_slice = self.slice_slider.value()
+                        else:
+                            self.current_slice = 0
+                            self.slice_slider.setValue(0)
                 else:
-                    self.slice_slider.setEnabled(True)
-                    self.slice_slider.setMinimum(0)
-                    self.slice_slider.setMaximum(self.channel_data[channel_index].shape[0] - 1)
-                    if self.slice_slider.value() < self.channel_data[channel_index].shape[0] - 1:
-                        self.current_slice = self.slice_slider.value()
-                    else:
-                        self.current_slice = 0
-                        self.slice_slider.setValue(0)
-            else:
-                self.slice_slider.setEnabled(False)
+                    self.slice_slider.setEnabled(False)
+            except:
+                pass
 
             
             # If this is the first channel loaded, make it active
@@ -3083,13 +3261,16 @@ class ImageViewerWindow(QMainWindow):
             self.min_max[channel_index][1] = np.max(self.channel_data[channel_index])
             self.volume_dict[channel_index] = None #reset volumes
 
-            if assign_shape: #keep original shape tracked to undo resampling.
-                if self.original_shape is None:
-                    self.original_shape = self.channel_data[channel_index].shape
-                elif self.original_shape[0] < self.channel_data[channel_index].shape[0] or self.original_shape[1] < self.channel_data[channel_index].shape[1] or self.original_shape[2] < self.channel_data[channel_index].shape[2]:
-                    self.original_shape = self.channel_data[channel_index].shape
-                if len(self.original_shape) == 4:
-                    self.original_shape = (self.original_shape[0], self.original_shape[1], self.original_shape[2])
+            try:
+                if assign_shape: #keep original shape tracked to undo resampling.
+                    if self.original_shape is None:
+                        self.original_shape = self.channel_data[channel_index].shape
+                    elif self.original_shape[0] < self.channel_data[channel_index].shape[0] or self.original_shape[1] < self.channel_data[channel_index].shape[1] or self.original_shape[2] < self.channel_data[channel_index].shape[2]:
+                        self.original_shape = self.channel_data[channel_index].shape
+                    if len(self.original_shape) == 4:
+                        self.original_shape = (self.original_shape[0], self.original_shape[1], self.original_shape[2])
+            except:
+                pass
 
             self.update_display(reset_resize = reset_resize)
 
@@ -3262,6 +3443,8 @@ class ImageViewerWindow(QMainWindow):
                 elif ch_index == 3:
                     my_network.save_id_overlay(filename=filename)
                 elif ch_index == 4:
+                    if self.mini_overlay == True:
+                        self.create_highlight_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
                     if filename == None:
                         filename = "Highlighted_Element.tif"
                     tifffile.imwrite(f"{filename}", self.highlight_overlay)
@@ -3306,7 +3489,11 @@ class ImageViewerWindow(QMainWindow):
         if self.pending_slice is not None:
             slice_value, view_settings = self.pending_slice
             self.current_slice = slice_value
+            if self.mini_overlay == True: #If we are rendering the highlight overlay for selected values one at a time.
+                self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
             self.update_display(preserve_zoom=view_settings)
+            if self.machine_window is not None:
+                self.machine_window.poke_segmenter()
             self.pending_slice = None
 
     def update_brightness(self, channel_index, values):
@@ -3442,7 +3629,15 @@ class ImageViewerWindow(QMainWindow):
             self.create_highlight_overlay_slice(self.targs, bounds = self.bounds)
 
         # Add highlight overlay if it exists
-        if self.highlight_overlay is not None and self.highlight and self.machine_window is None:
+        if self.mini_overlay and self.highlight and self.machine_window is None:
+            highlight_cmap = LinearSegmentedColormap.from_list(
+                'highlight',
+                [(0, 0, 0, 0), (1, 1, 0, 1)]  # yellow
+            )
+            self.ax.imshow(self.mini_overlay_data,
+                         cmap=highlight_cmap,
+                         alpha=0.5)
+        elif self.highlight_overlay is not None and self.highlight and self.machine_window is None:
             highlight_slice = self.highlight_overlay[self.current_slice]
             highlight_cmap = LinearSegmentedColormap.from_list(
                 'highlight',
@@ -3968,7 +4163,13 @@ class CustomTableView(QTableView):
                         # Navigate to the Z-slice
                         self.parent.slice_slider.setValue(int(centroid[0]))
                         print(f"Found node {value} at Z-slice {centroid[0]}")
-                        self.parent.create_highlight_overlay(node_indices=[value])
+                        if self.parent.channel_data[0].shape[0] * self.parent.channel_data[0].shape[1] * self.parent.channel_data[0].shape[2] > self.parent.mini_thresh:
+                            self.parent.mini_overlay = True
+                            self.parent.create_mini_overlay(node_indices = [value])
+                        else:
+                            self.parent.create_highlight_overlay(node_indices=[value])
+                        self.parent.clicked_values['nodes'] = []
+                        self.parent.clicked_values['edges'] = []
                         self.parent.clicked_values['nodes'].append(value)
                         
                         # Highlight the value in both tables if it exists
@@ -3992,7 +4193,13 @@ class CustomTableView(QTableView):
                         # Navigate to the Z-slice
                         self.parent.slice_slider.setValue(int(centroid[0]))
                         print(f"Found edge {value} at Z-slice {centroid[0]}")
-                        self.parent.create_highlight_overlay(edge_indices=[value])
+                        if self.parent.channel_data[1].shape[0] * self.parent.channel_data[1].shape[1] * self.parent.channel_data[1].shape[2] > self.parent.mini_thresh:
+                            self.parent.mini_overlay = True
+                            self.parent.create_mini_overlay(edge_indices = [value])
+                        else:
+                            self.parent.create_highlight_overlay(edge_indices=[value])
+                        self.parent.clicked_values['nodes'] = []
+                        self.parent.clicked_values['edges'] = []
                         self.parent.clicked_values['edges'].append(value)
 
                         # Highlight the value in both tables if it exists
@@ -4020,12 +4227,24 @@ class CustomTableView(QTableView):
                 self.parent.slice_slider.setValue(int(centroid1[0]))
                 print(f"Found node pair {value[0]} and {value[1]} at Z-slices {centroid1[0]} and {centroid2[0]}, respectively")
                 try:
-                    self.parent.create_highlight_overlay(node_indices=[int(value[0]), int(value[1])], edge_indices = int(value[2]))
+                    if self.parent.channel_data[0].shape[0] * self.parent.channel_data[0].shape[1] * self.parent.channel_data[0].shape[2] > self.parent.mini_thresh:
+                        self.parent.mini_overlay = True
+                        self.parent.create_mini_overlay(node_indices=[int(value[0]), int(value[1])], edge_indices = int(value[2]))
+                    else:
+                        self.parent.create_highlight_overlay(node_indices=[int(value[0]), int(value[1])], edge_indices = int(value[2]))
+                    self.parent.clicked_values['nodes'] = []
+                    self.parent.clicked_values['edges'] = []
                     self.parent.clicked_values['edges'].append(value[2])
                     self.parent.clicked_values['nodes'].append(value[0])
                     self.parent.clicked_values['nodes'].append(value[1])
                 except:
-                    self.parent.create_highlight_overlay(node_indices=[int(value[0]), int(value[1])])
+                    if self.parent.channel_data[0].shape[0] * self.parent.channel_data[0].shape[1] * self.parent.channel_data[0].shape[2] > self.parent.mini_thresh:
+                        self.parent.mini_overlay = True
+                        self.parent.create_mini_overlay(node_indices=[int(value[0]), int(value[1])])
+                    else:
+                        self.parent.create_highlight_overlay(node_indices=[int(value[0]), int(value[1])])
+                    self.parent.clicked_values['nodes'] = []
+                    self.parent.clicked_values['edges'] = []
                     self.parent.clicked_values['nodes'].append(value[0])
                     self.parent.clicked_values['nodes'].append(value[1])
 
@@ -4571,6 +4790,8 @@ class Show3dDialog(QDialog):
                             arrays_4d.append(channel)
 
             if self.parent().highlight_overlay is not None:
+                if self.parent().mini_overlay == True:
+                    self.parent().create_highlight_overlay(node_indices = self.parent().clicked_values['nodes'], edge_indices = self.parent().clicked_values['edges'])
                 arrays_3d.append(self.parent().highlight_overlay)
                 colors.append(color_template[4])
         
@@ -4705,6 +4926,45 @@ class IdOverlayDialog(QDialog):
 
         self.accept()
 
+class ColorOverlayDialog(QDialog):
+
+    def __init__(self, parent=None):
+
+        super().__init__(parent)
+        self.setWindowTitle("Generate Node (or Edge) -> Color Overlay?")
+        self.setModal(True)
+
+        layout = QFormLayout(self)
+
+        self.down_factor = QLineEdit("")
+        layout.addRow("down_factor (for speeding up overlay generation - optional):", self.down_factor)
+
+        # Add Run button
+        run_button = QPushButton("Generate (Will go to Overlay 2)")
+        run_button.clicked.connect(self.coloroverlay)
+        layout.addWidget(run_button)
+
+    def coloroverlay(self):
+
+        down_factor = float(self.down_factor.text()) if self.down_factor.text().strip() else None
+
+        if self.parent().active_channel == 0:
+            mode = 0
+            self.sort = 'Node'
+        else:
+            mode = 1
+            self.sort = 'Edge'
+
+
+        result, legend = my_network.node_to_color(down_factor = down_factor, mode = mode)
+
+        self.parent().format_for_upperright_table(legend, f'{self.sort} Id', f'Encoding Val: {self.sort}', 'Legend')
+
+
+        self.parent().load_channel(3, channel_data = result, data = True)
+
+        self.accept()
+
 
 class ShuffleDialog(QDialog):
 
@@ -4745,11 +5005,15 @@ class ShuffleDialog(QDialog):
             accepted_target = self.target_selector.currentIndex()
 
             if accepted_mode == 4:
+                if self.parent().mini_overlay == True:
+                    self.parent().create_highlight_overlay(node_indices = self.parent().clicked_values['nodes'], edge_indices = self.parent().clicked_values['edges'])
                 active_data = self.parent().highlight_overlay
             else:
                 active_data = self.parent().channel_data[accepted_mode]
 
             if accepted_target == 4:
+                if self.parent().mini_overlay == True:
+                    self.parent().create_highlight_overlay(node_indices = self.parent().clicked_values['nodes'], edge_indices = self.parent().clicked_values['edges'])
                 target_data = self.parent().highlight_overlay
             else:
                 target_data = self.parent().channel_data[accepted_target]
@@ -5952,7 +6216,18 @@ class MachineWindow(QMainWindow):
         self.GPU.setChecked(False)
         self.GPU.clicked.connect(self.toggle_GPU)
         self.use_gpu = False
+        self.two = QPushButton("Train By 2D Slice Patterns (Cheaper - CPU only)")
+        self.two.setCheckable(True)
+        self.two.setChecked(True)
+        self.two.clicked.connect(self.toggle_two)
+        self.use_two = True
+        self.three = QPushButton("Train by 3D Patterns")
+        self.three.setCheckable(True)
+        self.three.setChecked(False)
+        self.three.clicked.connect(self.toggle_three)
         processing_layout.addWidget(self.GPU)
+        processing_layout.addWidget(self.two)
+        processing_layout.addWidget(self.three)
         processing_group.setLayout(processing_layout)
 
         # Group 3: Training Options
@@ -5970,6 +6245,7 @@ class MachineWindow(QMainWindow):
         segmentation_group = QGroupBox("Segmentation")
         segmentation_layout = QVBoxLayout()
         seg_button = QPushButton("Preview Segment")
+        self.seg_button = seg_button
         seg_button.clicked.connect(self.start_segmentation)
         full_button = QPushButton("Segment All")
         full_button.clicked.connect(self.segment)
@@ -5987,9 +6263,12 @@ class MachineWindow(QMainWindow):
         self.setCentralWidget(main_widget)
 
         self.trained = False
+        self.previewing = False
 
 
         self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=True)
+        self.segmentation_worker = None
+
 
 
     def toggle_GPU(self):
@@ -5997,6 +6276,27 @@ class MachineWindow(QMainWindow):
 
         self.use_gpu = self.GPU.isChecked()
 
+    def toggle_two(self):
+        if self.two.isChecked():
+            # If button two is checked, ensure button three is unchecked
+            self.three.setChecked(False)
+            self.use_two = True
+        else:
+            # If button three is checked, ensure button two is unchecked
+            self.three.setChecked(True)
+            self.use_two = False
+
+    def toggle_three(self):
+        if self.three.isChecked():
+            # If button two is checked, ensure button three is unchecked
+            self.two.setChecked(False)
+            self.use_two = False
+        else:
+            # If button three is checked, ensure button two is unchecked
+            self.two.setChecked(True)
+            self.use_two = True
+
+
 
     def toggle_foreground(self):
 
@@ -6045,15 +6345,29 @@ class MachineWindow(QMainWindow):
         self.kill_segmentation()
         # Wait a bit for cleanup
         time.sleep(0.1)
-        self.segmenter.train_batch(self.parent().channel_data[2], speed = speed, use_gpu = self.use_gpu)
-        self.trained = True
+        if not self.use_two:
+            self.previewing = False
+        try:
+            self.segmenter.train_batch(self.parent().channel_data[2], speed = speed, use_gpu = self.use_gpu, use_two = self.use_two)
+            self.trained = True
+        except MemoryError:
+            QMessageBox.critical(
+                self,
+                "Alert",
+                "Out of memory computing feature maps. Note these for 3D require 7x the RAM of the active image (or 9x for the detailed map).\n Please use 2D slice models if you do not have enough RAM."
+            )
+
 
     def start_segmentation(self):
 
         self.kill_segmentation()
         time.sleep(0.1)
 
-        print("Beginning new segmentation...")
+        if self.use_two:
+            self.previewing = True
+        else:
+            print("Beginning new segmentation...")
+
 
         if self.parent().active_channel == 0:
             if self.parent().channel_data[0] is not None:
@@ -6067,7 +6381,7 @@ class MachineWindow(QMainWindow):
         if not self.trained:
             return
         else:
-            self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu)
+            self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu, self.use_two, self.previewing, self)
             self.segmentation_worker.chunk_processed.connect(self.update_display)  # Just update display
             self.segmentation_worker.finished.connect(self.segmentation_finished)
             current_xlim = self.parent().ax.get_xlim()
@@ -6095,10 +6409,42 @@ class MachineWindow(QMainWindow):
         msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
         return msg.exec() == QMessageBox.StandardButton.Yes
 
+
+
+    def check_for_z_change(self):
+        current_z = self.parent().current_slice
+        if not hasattr(self, '_last_z'):
+            self._last_z = current_z
+            return False
+        
+        changed = (self._last_z != current_z)
+        self._last_z = current_z
+        
+        if changed and self.previewing and self.segmentation_worker is not None:
+            self.segmentation_worker.stop()
+            time.sleep(0.1)
+            
+            # Force regeneration of chunks
+            self.segmenter.realtimechunks = None
+            
+            # Restart the worker
+            self.start_segmentation()
+            
+        return changed
+
     def update_display(self):
         if not hasattr(self, '_last_update'):
             self._last_update = 0
         
+        current_z = self.parent().current_slice
+        if not hasattr(self, '_last_z'):
+            self._last_z = current_z
+
+        self._last_z = current_z
+
+        if self.previewing:
+            changed = self.check_for_z_change()
+
         current_time = time.time()
         if current_time - self._last_update >= 1:  # Match worker's interval
             try:
@@ -6116,19 +6462,84 @@ class MachineWindow(QMainWindow):
             except Exception as e:
                 print(f"Display update error: {e}")
 
+    def poke_segmenter(self):
+        if self.use_two and self.previewing:
+            try:
+                # Clear any processing flags in the segmenter
+                if hasattr(self.segmenter, '_currently_processing'):
+                    self.segmenter._currently_processing = None
+                    
+                # Force regenerating the worker
+                if self.segmentation_worker is not None:
+                    self.kill_segmentation()
+                    
+                time.sleep(0.2)
+                self.start_segmentation()
+                
+            except Exception as e:
+                print(f"Error in poke_segmenter: {e}")
+                import traceback
+                traceback.print_exc()
+
     def segmentation_finished(self):
-        print("Segmentation completed")
+        if not self.use_two:
+            print("Segmentation completed")
+        
         current_xlim = self.parent().ax.get_xlim()
         current_ylim = self.parent().ax.get_ylim()
         self.parent().update_display(preserve_zoom=(current_xlim, current_ylim))
+        
+        # Store the current z position before killing the worker
+        current_z = self.parent().current_slice
+        
+        # Clean up the worker
         self.kill_segmentation()
+        self.segmentation_worker = None
         time.sleep(0.1)
+        
+        # Auto-restart for 2D preview mode only if certain conditions are met
+        if self.previewing and self.use_two:
+            # Track when this slice was last processed
+            if not hasattr(self, '_processed_slices'):
+                self._processed_slices = {}
+                
+            current_time = time.time()
+            
+            # Check if we've recently tried to process this slice (to prevent loops)
+            recently_processed = False
+            if current_z in self._processed_slices:
+                time_since_last_attempt = current_time - self._processed_slices[current_z]
+                recently_processed = time_since_last_attempt < 5.0  # 5 second cooldown
+                
+            if not recently_processed:
+                self._processed_slices[current_z] = current_time
+                
+                # Reset any processing flags in the segmenter
+                if hasattr(self.segmenter, '_currently_processing'):
+                    self.segmenter._currently_processing = None
+                    
+                # Create a new worker after a brief delay
+                QTimer.singleShot(500, self.start_segmentation)
+
 
 
     def kill_segmentation(self):
-        if hasattr(self, 'segmentation_worker'):
+        if hasattr(self, 'segmentation_worker') and self.segmentation_worker is not None:
+            # Signal the thread to stop
             self.segmentation_worker.stop()
-            del self.segmentation_worker  # Clean up reference
+            
+            # Wait for the thread to finish
+            if self.segmentation_worker.isRunning():
+                self.segmentation_worker.wait(1000)  # Wait up to 1 second
+                
+                # If thread is still running after timeout, try to force termination
+                if self.segmentation_worker.isRunning():
+                    self.segmentation_worker.terminate()
+                    self.segmentation_worker.wait()  # Wait for it to be terminated
+            
+            # Now safe to delete
+            del self.segmentation_worker
+            self.segmentation_worker = None
 
 
     def segment(self):
@@ -6141,6 +6552,8 @@ class MachineWindow(QMainWindow):
             self.kill_segmentation()
             time.sleep(0.1)
 
+            self.previewing = False
+
             if self.parent().active_channel == 0:
                 if self.parent().channel_data[0] is not None:
                     active_data = self.parent().channel_data[0]
@@ -6170,20 +6583,26 @@ class MachineWindow(QMainWindow):
 
         self.parent().update_display()
 
+        self.previewing = False
+
         print("Finished segmentation moved to Overlay 2. Use File -> Save(As) for disk saving.")
 
     def closeEvent(self, event):
         if self.parent().isVisible():
             if self.confirm_close_dialog():
-
+                # Clean up resources before closing
                 if self.brush_button.isChecked():
                     self.silence_button()
                     self.toggle_brush_mode()
+                
                 self.parent().pen_button.setEnabled(True)
                 self.parent().brush_mode = False
-                self.parent().machine_window = None
+                
+                # Kill the segmentation thread and wait for it to finish
                 self.kill_segmentation()
-                time.sleep(0.1)
+                time.sleep(0.2)  # Give additional time for cleanup
+                
+                self.parent().machine_window = None
             else:
                 event.ignore()
 
@@ -6194,48 +6613,84 @@ class SegmentationWorker(QThread):
     finished = pyqtSignal()
     chunk_processed = pyqtSignal()
     
-    def __init__(self, highlight_overlay, segmenter, use_gpu):
+    def __init__(self, highlight_overlay, segmenter, use_gpu, use_two, previewing, machine_window):
         super().__init__()
         self.overlay = highlight_overlay
         self.segmenter = segmenter
         self.use_gpu = use_gpu
+        self.use_two = use_two
+        self.previewing = previewing
+        self.machine_window = machine_window
         self._stop = False
         self.update_interval = 1  # Increased to 500ms
         self.chunks_since_update = 0
         self.chunks_per_update = 5  # Only update every 5 chunks
+        self.poked = False # If it should wake up or not
         self.last_update = time.time()
         
     def stop(self):
         self._stop = True
+
+    def get_poked(self):
+        self.poked = True
         
     def run(self):
         try:
             self.overlay.fill(False)
             
-            for foreground_coords, background_coords in self.segmenter.segment_volume_realtime(gpu = self.use_gpu):
-                if self._stop:
-                    break
+            # Remember the starting z position
+            self.starting_z = self.segmenter.current_z
+            
+            if self.previewing and self.use_two:
+                # Process current z-slice in chunks
+                current_z = self.segmenter.current_z
+                
+                # Process the slice with chunked generator
+                for foreground, background in self.segmenter.segment_slice_chunked(current_z):
+                    if self._stop:
+                        break
                     
-                for z,y,x in foreground_coords:
-                    self.overlay[z,y,x] = 1
-                for z,y,x in background_coords:
-                    self.overlay[z,y,x] = 2
-                
-                # Update only after several chunks AND minimum time interval
-                self.chunks_since_update += 1
-                current_time = time.time()
-                if (self.chunks_since_update >= self.chunks_per_update and 
-                    current_time - self.last_update >= self.update_interval):
-                    self.chunk_processed.emit()
-                    self.chunks_since_update = 0
-                    self.last_update = current_time
-
+                    # Update the overlay
+                    for z,y,x in foreground:
+                        self.overlay[z,y,x] = 1
+                    for z,y,x in background:
+                        self.overlay[z,y,x] = 2
+                    
+                    # Signal update after each chunk
+                    self.chunks_since_update += 1
+                    current_time = time.time()
+                    if (self.chunks_since_update >= self.chunks_per_update and 
+                        current_time - self.last_update >= self.update_interval):
+                        self.chunk_processed.emit()
+                        self.chunks_since_update = 0
+                        self.last_update = current_time
+                
+            else:
+                # Original 3D approach
+                for foreground_coords, background_coords in self.segmenter.segment_volume_realtime(gpu=self.use_gpu):
+                    if self._stop:
+                        break
+                    
+                    for z,y,x in foreground_coords:
+                        self.overlay[z,y,x] = 1
+                    for z,y,x in background_coords:
+                        self.overlay[z,y,x] = 2
+                    
+                    self.chunks_since_update += 1
+                    current_time = time.time()
+                    if (self.chunks_since_update >= self.chunks_per_update and 
+                        current_time - self.last_update >= self.update_interval):
+                        self.chunk_processed.emit()
+                        self.chunks_since_update = 0
+                        self.last_update = current_time
+            
             self.finished.emit()
             
         except Exception as e:
             print(f"Error in segmentation: {e}")
-            raise
-        
+            import traceback
+            traceback.print_exc()
+
     def run_batch(self):
         try:
             foreground_coords, _ = self.segmenter.segment_volume()
@@ -6851,6 +7306,8 @@ class MaskDialog(QDialog):
             output_target = self.output_selector.currentIndex()
 
             if accepted_mode == 4:
+                if self.parent().mini_overlay == True:
+                    self.parent().create_highlight_overlay(node_indices = self.parent().clicked_values['nodes'], edge_indices = self.parent().clicked_values['edges'])
                 active_data = self.parent().highlight_overlay
             else:
                 active_data = self.parent().channel_data[accepted_mode]