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

nettracer3d/nettracer_gui.py

@@ -397,6 +397,8 @@ 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_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 +445,8 @@ 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.
+
         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 +467,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 = []
@@ -628,6 +635,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 binary overlay highlighting specific nodes and/or edges using parallel processing, one slice at a time for efficiency.
+        
+        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:
+            self.highlight_overlay = None
+            self.update_display(preserve_zoom=(current_xlim, current_ylim))
+            return
+            
+        # Get the shape of the full array from any existing channel
+        for channel in self.channel_data:
+            if channel is not None:
+                full_shape = channel.shape
+                break
+        else:
+            return  # No valid channels to get shape from
+            
+        # Initialize full-size overlay
+        if self.highlight_overlay is None:
+            self.highlight_overlay = 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[1] // 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[1], chunk_size):
+                end = min(i + chunk_size, array_shape[1])
+                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
+        self.highlight_overlay[self.current_slice, :, :] = np.zeros_like(self.highlight_overlay[self.current_slice, :, :])
+        if node_overlay is not None:
+            self.highlight_overlay[self.current_slice, :, :] = np.maximum(self.highlight_overlay[self.current_slice, :, :], node_overlay)
+        if edge_overlay is not None:
+            self.highlight_overlay[self.current_slice, :, :] = np.maximum(self.highlight_overlay[self.current_slice, :, :], edge_overlay)
+
+                
+        # Update display
+        self.update_display(preserve_zoom=(current_xlim, current_ylim))
+
 
 
         
@@ -890,14 +1001,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 +1092,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 +1172,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 +1233,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}")
@@ -1286,7 +1430,7 @@ class ImageViewerWindow(QMainWindow):
                 self.load_channel(1, my_network.edges, True)
             self.highlight_overlay = None
             self.update_display()
-            print("Network is not updated automatically, please recompute if necesarry. Identities are not automatically updated.")
+            print("Network is not updated automatically, please recompute if necessary. Identities are not automatically updated.")
             self.show_centroid_dialog()
 
         except Exception as e:
@@ -1315,7 +1459,7 @@ class ImageViewerWindow(QMainWindow):
 
 
             if len(self.clicked_values['edges']) > 0:
-                self.create_highlight_overlay(node_indices = self.clicked_values['edges'])
+                self.create_highlight_overlay(edge_indices = self.clicked_values['edges'])
                 mask = self.highlight_overlay == 0
                 my_network.edges = my_network.edges * mask
                 self.load_channel(1, my_network.edges, True)
@@ -1393,8 +1537,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 +2033,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 +2049,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 +2294,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 +2456,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 +2738,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)
@@ -2954,14 +3119,46 @@ class ImageViewerWindow(QMainWindow):
 
         try:
             if not data:  # For solo loading
-                import tifffile
                 filename, _ = QFileDialog.getOpenFileName(
                     self,
                     f"Load Channel {channel_index + 1}",
                     "",
-                    "TIFF Files (*.tif *.tiff)"
+                    "Image Files (*.tif *.tiff *.nii *.jpg *.jpeg *.png)"
                 )
-                self.channel_data[channel_index] = tifffile.imread(filename)
+                
+                if not filename:
+                    return
+                    
+                file_extension = filename.lower().split('.')[-1]
+                
+                try:
+                    if file_extension in ['tif', 'tiff']:
+                        import tifffile
+                        self.channel_data[channel_index] = tifffile.imread(filename)
+                        
+                    elif file_extension == 'nii':
+                        import nibabel as nib
+                        nii_img = nib.load(filename)
+                        # Get data and transpose to match TIFF orientation
+                        # If X needs to become Z, we move axis 2 (X) to position 0 (Z)
+                        data = nii_img.get_fdata()
+                        self.channel_data[channel_index] = np.transpose(data, (2, 1, 0))
+                        
+                    elif file_extension in ['jpg', 'jpeg', 'png']:
+                        from PIL import Image
+                        
+                        with Image.open(filename) as img:
+                            # Convert directly to numpy array, keeping color if present
+                            self.channel_data[channel_index] = np.array(img)
+                            
+                            # Debug info to check shape
+                            print(f"Loaded image shape: {self.channel_data[channel_index].shape}")
+                            
+                except ImportError as e:
+                    QMessageBox.critical(self, "Error", f"Required library not installed: {str(e)}")
+                except Exception as e:
+                    QMessageBox.critical(self, "Error", f"Error loading image: {str(e)}")
+
 
                 if len(self.channel_data[channel_index].shape) == 2:  # handle 2d data
                     self.channel_data[channel_index] = np.expand_dims(self.channel_data[channel_index], axis=0)
@@ -2969,24 +3166,33 @@ 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
 
             for i in range(4): #Try to ensure users don't load in different sized arrays
                 if self.channel_data[i] is None or i == channel_index or data:
                     if self.highlight_overlay is not None: #Make sure highlight overlay is always the same shape as new images
-                        if self.channel_data[i].shape[:3] != self.highlight_overlay.shape:
-                            self.resizing = True
-                            reset_resize = True
-                            self.highlight_overlay = None
+                        try:
+                            if self.channel_data[i].shape[:3] != self.highlight_overlay.shape:
+                                self.resizing = True
+                                reset_resize = True
+                                self.highlight_overlay = None
+                        except:
+                            pass
                     continue
                 else:
                     old_shape = self.channel_data[i].shape[:3] #Ask user to resize images that are shaped differently
@@ -3015,27 +3221,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
@@ -3048,19 +3257,24 @@ 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)
 
 
                 
         except Exception as e:
+            import traceback
+            print(traceback.format_exc())
             if not data:
                 from PyQt6.QtWidgets import QMessageBox
                 QMessageBox.critical(
@@ -3225,6 +3439,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)
@@ -3269,6 +3485,8 @@ 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)
             self.pending_slice = None
 
@@ -3931,7 +4149,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
@@ -3955,7 +4179,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
@@ -3983,12 +4213,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])
 
@@ -4534,6 +4776,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])
         
@@ -4668,6 +4912,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):
 
@@ -4708,11 +4991,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]
@@ -6814,6 +7101,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]
@@ -6916,11 +7205,16 @@ class WatershedDialog(QDialog):
         self.directory.setPlaceholderText("Leave empty for None")
         layout.addRow("Output Directory:", self.directory)
 
-        active_shape = self.parent().channel_data[self.parent().active_channel].shape[0]
+        try:
 
-        if active_shape == 1:
-            self.default = 0.2
-        else:
+            active_shape = self.parent().channel_data[self.parent().active_channel].shape[0]
+
+            if active_shape == 1:
+                self.default = 0.2
+            else:
+                self.default = 0.05
+
+        except:
             self.default = 0.05
 
         
@@ -6949,6 +7243,9 @@ class WatershedDialog(QDialog):
         self.predownsample2.setPlaceholderText("Leave empty for None")
         layout.addRow("Smart Label GPU Downsample:", self.predownsample2)
         
+        layout.addRow("Note:", QLabel(f"If the optimal proportion watershed output is still labeling spatially seperated objects with the same label, try right placing the result in nodes or edges\nthen right click the image and choose 'select all', followed by right clicking and 'selection' -> 'split non-touching labels'."))
+
+
         # Add Run button
         run_button = QPushButton("Run Watershed")
         run_button.clicked.connect(self.run_watershed)