nettracer3d 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

nettracer3d/nettracer_gui.py

@@ -22,7 +22,7 @@ from PyQt6.QtGui import (QFont, QCursor, QColor, QPixmap, QFontMetrics, QPainter
 import tifffile
 import copy
 import multiprocessing as mp
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from functools import partial
 from nettracer3d import segmenter
 try:
@@ -33,6 +33,9 @@ from nettracer3d import excelotron
 import threading
 import queue
 from threading import Lock
+from scipy import ndimage
+import os
+from . import painting
 
 
 
@@ -208,7 +211,7 @@ class ImageViewerWindow(QMainWindow):
         buttons_widget = QWidget()
         buttons_layout = QHBoxLayout(buttons_widget)
 
-        # Create zoom button
+        # "Create" zoom button
         self.zoom_button = QPushButton("🔍")
         self.zoom_button.setCheckable(True)
         self.zoom_button.setFixedSize(40, 40)
@@ -291,6 +294,14 @@ class ImageViewerWindow(QMainWindow):
             
             control_layout.addWidget(channel_container)
 
+        self.show_channels = QPushButton("✓")
+        self.show_channels.setCheckable(True)
+        self.show_channels.setChecked(True)
+        self.show_channels.setFixedSize(20, 20)
+        self.show_channels.clicked.connect(self.toggle_chan_load)
+        control_layout.addWidget(self.show_channels)
+        self.chan_load = True
+
         # Create the main widget and layout
         main_widget = QWidget()
         self.setCentralWidget(main_widget)
@@ -386,7 +397,7 @@ class ImageViewerWindow(QMainWindow):
         # Create both table views
         self.network_table = CustomTableView(self)
         self.selection_table = CustomTableView(self)
-        empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+        empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
         self.selection_table.setModel(PandasModel(empty_df))
         self.network_table.setAlternatingRowColors(True)
         self.selection_table.setAlternatingRowColors(True)
@@ -444,18 +455,8 @@ class ImageViewerWindow(QMainWindow):
         self.excel_manager.data_received.connect(self.handle_excel_data)
         self.prev_coms = None
 
-        self.paint_timer = QTimer()
-        self.paint_timer.timeout.connect(self.flush_paint_updates)
-        self.paint_timer.setSingleShot(True)
-        self.pending_paint_update = False
         self.static_background = None
         
-        # Threading for paint operations
-        self.paint_queue = queue.Queue()
-        self.paint_lock = Lock()
-        self.paint_worker = threading.Thread(target=self.paint_worker_loop, daemon=True)
-        self.paint_worker.start()
-        
         # Background caching for blitting
         self.paint_session_active = False
         
@@ -1773,7 +1774,7 @@ class ImageViewerWindow(QMainWindow):
                 my_network.network_lists = my_network.network_lists
 
                 if not hasattr(my_network, 'network_lists') or my_network.network_lists is None:
-                    empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+                    empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
                     model = PandasModel(empty_df)
                     self.network_table.setModel(model)
                 else:
@@ -1795,58 +1796,149 @@ class ImageViewerWindow(QMainWindow):
         except Exception as e:
             print(f"An error has occured: {e}")
 
-    def separate_nontouching_objects(self, input_array, max_val=0):
+
+    def process_single_label_bbox(args):
         """
-        optimized version using advanced indexing.
+        Worker function to process a single label within its bounding box
+        This function will run in parallel
         """
+        label_subarray, original_label, bbox_slices, start_new_label = args
+        
+        try:
+            # Create binary mask for this label only
+            binary_mask = label_subarray == original_label
+            
+            if not np.any(binary_mask):
+                return None, start_new_label, bbox_slices
+            
+            # Find connected components in the subarray
+            labeled_cc, num_cc = n3d.label_objects(binary_mask)
+            
+            if num_cc == 0:
+                return None, start_new_label, bbox_slices
+            
+            # Create output subarray with new labels
+            output_subarray = np.zeros_like(label_subarray)
+            
+            # Assign new consecutive labels starting from start_new_label
+            for cc_id in range(1, num_cc + 1):
+                cc_mask = labeled_cc == cc_id
+                new_label = start_new_label + cc_id - 1
+                output_subarray[cc_mask] = new_label
+            
+            # Return the processed subarray, number of components created, and bbox info
+            return output_subarray, start_new_label + num_cc, bbox_slices
+            
+        except Exception as e:
+            print(f"Error processing label {original_label}: {e}")
+            return None, start_new_label, bbox_slices
 
+    def separate_nontouching_objects(self, input_array, max_val=0):
+        """
+        Ultra-optimized version using find_objects directly without remapping
+        """
         print("Splitting nontouching objects")
-
+        
         binary_mask = input_array > 0
-        labeled_array, _ = n3d.label_objects(binary_mask)
+        if not np.any(binary_mask):
+            return np.zeros_like(input_array)
         
-        # Create a compound key for each (original_label, connected_component) pair
-        # This avoids the need for explicit mapping
-        mask = binary_mask
-        compound_key = input_array[mask] * (labeled_array.max() + 1) + labeled_array[mask]
+        unique_labels = np.unique(input_array[binary_mask])
+        print(f"Processing {len(unique_labels)} unique labels")
         
-        # Get unique compound keys and create new labels
-        unique_keys, inverse_indices = np.unique(compound_key, return_inverse=True)
-        new_labels = np.arange(max_val + 1, max_val + 1 + len(unique_keys))
+        # Get all bounding boxes at once - this is very fast
+        bounding_boxes = ndimage.find_objects(input_array)
         
-        # Create output array
-        output_array = np.zeros_like(input_array)
-        output_array[mask] = new_labels[inverse_indices]
+        # Prepare work items - just check if bounding box exists for each label
+        work_items = []
+        for orig_label in unique_labels:
+            # find_objects returns list where index = label - 1
+            bbox_index = orig_label - 1
+            
+            if (bbox_index >= 0 and 
+                bbox_index < len(bounding_boxes) and 
+                bounding_boxes[bbox_index] is not None):
+                
+                bbox = bounding_boxes[bbox_index]
+                work_items.append((orig_label, bbox))
+        
+        print(f"Created {len(work_items)} work items")
+        
+        # If we have work items, process them
+        if len(work_items) == 0:
+            print("No valid work items found!")
+            return np.zeros_like(input_array)
+        
+        def process_label_minimal(item):
+            orig_label, bbox = item
+            try:
+                subarray = input_array[bbox]
+                binary_sub = subarray == orig_label
+                
+                if not np.any(binary_sub):
+                    return orig_label, bbox, None, 0
+                    
+                labeled_sub, num_cc = n3d.label_objects(binary_sub)
+                return orig_label, bbox, labeled_sub, num_cc
+                
+            except Exception as e:
+                print(f"Error processing label {orig_label}: {e}")
+                return orig_label, bbox, None, 0
+        
+        # Execute in parallel
+        max_workers = min(mp.cpu_count(), len(work_items))
+        
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            results = list(executor.map(process_label_minimal, work_items))
         
+        # Reconstruct output array
+        output_array = np.zeros_like(input_array)
+        current_label = max_val + 1
+        total_components = 0
+        
+        for orig_label, bbox, labeled_sub, num_cc in results:
+            if num_cc > 0 and labeled_sub is not None:
+                print(f"Label {orig_label}: {num_cc} components")
+                # Remap labels and place in output
+                for cc_id in range(1, num_cc + 1):
+                    mask = labeled_sub == cc_id
+                    output_array[bbox][mask] = current_label
+                    current_label += 1
+                    total_components += 1
+        
+        print(f"Total components created: {total_components}")
         return output_array
 
     def handle_seperate(self):
-
+        """
+        Fixed version with proper mask handling and debugging
+        """
         try:
             # Handle nodes
             if len(self.clicked_values['nodes']) > 0:
-                self.create_highlight_overlay(node_indices=self.clicked_values['nodes'])
-                
-                # Create a boolean mask for highlighted values
-                self.highlight_overlay = self.highlight_overlay != 0
                 
-                # Create array with just the highlighted values
-                highlighted_nodes = self.highlight_overlay * my_network.nodes
+                # Create highlight overlay (this should preserve original label values)
+                self.create_highlight_overlay(node_indices=self.clicked_values['nodes'])
+                                
+                # DON'T convert to boolean yet - we need the original labels!
+                # Create a boolean mask for where we have highlighted values
+                highlight_mask = self.highlight_overlay != 0
                 
+                # Create array with just the highlighted values (preserving original labels)
+                highlighted_nodes = np.where(highlight_mask, my_network.nodes, 0)
+                                
                 # Get non-highlighted part of the array
-                non_highlighted = my_network.nodes * (~self.highlight_overlay)
-
-                if (highlighted_nodes==non_highlighted).all():
-                    max_val = 0
-                else:
-                    max_val = np.max(non_highlighted)
+                non_highlighted = np.where(highlight_mask, 0, my_network.nodes)
                 
+                # Calculate max_val properly
+                max_val = np.max(non_highlighted) if np.any(non_highlighted) else 0
+                                
                 # Process highlighted part
                 processed_highlights = self.separate_nontouching_objects(highlighted_nodes, max_val)
-                
+                                
                 # Combine back with non-highlighted parts
                 my_network.nodes = non_highlighted + processed_highlights
-                
+                                
                 self.load_channel(0, my_network.nodes, True)
             
             # Handle edges
@@ -1855,18 +1947,15 @@ class ImageViewerWindow(QMainWindow):
                 self.create_highlight_overlay(edge_indices=self.clicked_values['edges'])
                 
                 # Create a boolean mask for highlighted values
-                self.highlight_overlay = self.highlight_overlay != 0
+                highlight_mask = self.highlight_overlay != 0
                 
                 # Create array with just the highlighted values
-                highlighted_edges = self.highlight_overlay * my_network.edges
+                highlighted_edges = np.where(highlight_mask, my_network.edges, 0)
                 
                 # Get non-highlighted part of the array
-                non_highlighted = my_network.edges * (~self.highlight_overlay)
+                non_highlighted = np.where(highlight_mask, 0, my_network.edges)
 
-                if (highlighted_edges==non_highlighted).all():
-                    max_val = 0
-                else:
-                    max_val = np.max(non_highlighted)
+                max_val = np.max(non_highlighted) if np.any(non_highlighted) else 0
                 
                 # Process highlighted part
                 processed_highlights = self.separate_nontouching_objects(highlighted_edges, max_val)
@@ -1920,7 +2009,7 @@ class ImageViewerWindow(QMainWindow):
 
 
             if not hasattr(my_network, 'network_lists') or my_network.network_lists is None:
-                empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+                empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
                 model = PandasModel(empty_df)
                 self.network_table.setModel(model)
             else:
@@ -1961,7 +2050,7 @@ class ImageViewerWindow(QMainWindow):
             
             # Update the table
             if not hasattr(my_network, 'network_lists') or my_network.network_lists is None:
-                empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+                empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
                 model = PandasModel(empty_df)
                 self.network_table.setModel(model)
             else:
@@ -1993,7 +2082,7 @@ class ImageViewerWindow(QMainWindow):
             my_network.network_lists = my_network.network_lists
 
             if not hasattr(my_network, 'network_lists') or my_network.network_lists is None:
-                empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+                empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
                 model = PandasModel(empty_df)
                 self.network_table.setModel(model)
             else:
@@ -2044,6 +2133,12 @@ class ImageViewerWindow(QMainWindow):
             print(f"Error: {e}")
 
 
+    def toggle_chan_load(self):
+
+        if self.show_channels.isChecked():
+            self.chan_load = True
+        else:
+            self.chan_load = False
 
     def toggle_highlight(self):
         self.highlight = self.high_button.isChecked()
@@ -2066,11 +2161,6 @@ class ImageViewerWindow(QMainWindow):
         if self.zoom_mode:
             self.pan_button.setChecked(False)
 
-            if self.pan_mode or self.brush_mode:
-                current_xlim = self.ax.get_xlim()
-                current_ylim = self.ax.get_ylim()
-                self.update_display(preserve_zoom=(current_xlim, current_ylim))
-
             self.pen_button.setChecked(False)
             self.pan_mode = False
             self.brush_mode = False
@@ -2080,6 +2170,11 @@ class ImageViewerWindow(QMainWindow):
             if self.machine_window is not None:
                 self.machine_window.silence_button()
             self.canvas.setCursor(Qt.CursorShape.CrossCursor)
+            if self.pan_mode or self.brush_mode:
+                current_xlim = self.ax.get_xlim()
+                current_ylim = self.ax.get_ylim()
+                self.update_display(preserve_zoom=(current_xlim, current_ylim))
+
         else:
             if self.machine_window is None:
                 self.canvas.setCursor(Qt.CursorShape.ArrowCursor)
@@ -2091,10 +2186,6 @@ class ImageViewerWindow(QMainWindow):
         """Toggle pan mode on/off."""
         self.pan_mode = self.pan_button.isChecked()
         if self.pan_mode:
-            if self.brush_mode:
-                current_xlim = self.ax.get_xlim()
-                current_ylim = self.ax.get_ylim()
-                self.update_display(preserve_zoom=(current_xlim, current_ylim))
 
             self.zoom_button.setChecked(False)
             self.pen_button.setChecked(False)
@@ -2106,6 +2197,10 @@ class ImageViewerWindow(QMainWindow):
             if self.machine_window is not None:
                 self.machine_window.silence_button()
             self.canvas.setCursor(Qt.CursorShape.OpenHandCursor)
+            if self.brush_mode:
+                current_xlim = self.ax.get_xlim()
+                current_ylim = self.ax.get_ylim()
+                self.update_display(preserve_zoom=(current_xlim, current_ylim))
         else:
             current_xlim = self.ax.get_xlim()
             current_ylim = self.ax.get_ylim()
@@ -2120,6 +2215,20 @@ class ImageViewerWindow(QMainWindow):
         self.brush_mode = self.pen_button.isChecked()
         if self.brush_mode:
 
+            self.pm = painting.PaintManager(parent = self)
+
+            # Start virtual paint session
+            # Get current zoom to preserve it
+            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 self.pen_button.isChecked():
+                channel = self.active_channel
+            else:
+                channel = 2
+
+            self.pm.initiate_paint_session(channel, current_xlim, current_ylim)
+
             if self.pan_mode:
                 current_xlim = self.ax.get_xlim()
                 current_ylim = self.ax.get_ylim()
@@ -2339,37 +2448,6 @@ class ImageViewerWindow(QMainWindow):
         
         painter.end()
 
-    def get_line_points(self, x0, y0, x1, y1):
-        """Get all points in a line between (x0,y0) and (x1,y1) using Bresenham's algorithm."""
-        points = []
-        dx = abs(x1 - x0)
-        dy = abs(y1 - y0)
-        x, y = x0, y0
-        sx = 1 if x0 < x1 else -1
-        sy = 1 if y0 < y1 else -1
-        
-        if dx > dy:
-            err = dx / 2.0
-            while x != x1:
-                points.append((x, y))
-                err -= dy
-                if err < 0:
-                    y += sy
-                    err += dx
-                x += sx
-        else:
-            err = dy / 2.0
-            while y != y1:
-                points.append((x, y))
-                err -= dx
-                if err < 0:
-                    x += sx
-                    err += dy
-                y += sy
-                
-        points.append((x, y))
-        return points
-
     def get_current_mouse_position(self):
         # Get the main application's current mouse position
         cursor_pos = QCursor.pos()
@@ -2382,15 +2460,25 @@ class ImageViewerWindow(QMainWindow):
                 0 <= canvas_pos.y() < self.canvas.height()):
             return 0, 0  # Mouse is outside of the matplotlib canvas
         
-        # Convert from canvas widget coordinates to matplotlib data coordinates
-        x = canvas_pos.x()
-        y = canvas_pos.y()
-        
-        # Transform display coordinates to data coordinates
-        inv = self.ax.transData.inverted()
-        data_coords = inv.transform((x, y))
-        
-        return data_coords[0], data_coords[1]
+        # OPTION 1: Use matplotlib's built-in coordinate conversion
+        # This accounts for figure margins, subplot positioning, etc.
+        try:
+            # Get the figure and axes bounds
+            bbox = self.ax.bbox
+            
+            # Convert widget coordinates to figure coordinates
+            fig_x = canvas_pos.x()
+            fig_y = self.canvas.height() - canvas_pos.y()  # Flip Y coordinate
+            
+            # Check if within axes bounds
+            if (bbox.x0 <= fig_x <= bbox.x1 and bbox.y0 <= fig_y <= bbox.y1):
+                # Transform to data coordinates
+                data_coords = self.ax.transData.inverted().transform((fig_x, fig_y))
+                return data_coords[0], data_coords[1]
+            else:
+                return 0, 0
+        except:
+            pass
 
     def on_mouse_press(self, event):
         """Handle mouse press events."""
@@ -2429,61 +2517,50 @@ class ImageViewerWindow(QMainWindow):
 
 
         elif self.brush_mode:
+            """Handle brush mode with virtual painting."""
             if event.inaxes != self.ax:
                 return
             
             if event.button == 1 or event.button == 3:
+                if self.machine_window is not None:
+                    if self.machine_window.segmentation_worker is not None:
+                        self.machine_window.segmentation_worker.pause()
+                        
                 x, y = int(event.xdata), int(event.ydata)
-                # Get current zoom to preserve it
-
+                
+                # Get current zoom to preserve it 
                 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 event.button == 1 and self.can:
-                    self.update_display(preserve_zoom = (current_xlim, current_ylim))
+                
+                if event.button == 1 and getattr(self, 'can', False):
+                    self.update_display(preserve_zoom=(current_xlim, current_ylim), skip_paint_reinit = True)
                     self.handle_can(x, y)
                     return
-
+                    
+                # Determine erase mode and foreground/background
                 if event.button == 3:
                     self.erase = True
-                    self.update_display(preserve_zoom = (current_xlim, current_ylim))
                 else:
                     self.erase = False
-
-                self.painting = True
-                self.last_paint_pos = (x, y)
-
+                
+                # Determine foreground/background for machine window mode
+                foreground = getattr(self, 'foreground', True)
+                    
+                self.last_virtual_pos = (x, y)
+                
                 if self.pen_button.isChecked():
                     channel = self.active_channel
                 else:
                     channel = 2
 
-                # Paint at initial position
-                self.paint_at_position(x, y, self.erase, channel)
-                                
-                self.canvas.draw()
-
-                self.restore_channels = []
-                if not self.channel_visible[channel]:
-                    self.channel_visible[channel] = True
-
-                # No need to hide other channels or track restore_channels
-                self.restore_channels = []
-
-                if self.static_background is None:
-                    if self.machine_window is not None:
-                        self.update_display(preserve_zoom = (current_xlim, current_ylim))
-                    elif not self.erase:
-                        self.temp_chan = channel
-                        self.channel_data[4] = self.channel_data[channel]
-                        self.min_max[4] = copy.deepcopy(self.min_max[channel])
-                        self.channel_brightness[4] = copy.deepcopy(self.channel_brightness[channel])
-                        self.load_channel(channel, np.zeros_like(self.channel_data[channel]), data = True, preserve_zoom = (current_xlim, current_ylim), begin_paint = True)
-                        self.channel_visible[4] = True
-                    self.static_background = self.canvas.copy_from_bbox(self.ax.bbox)    
-
-                self.update_display_slice_optimized(channel, preserve_zoom=(current_xlim, current_ylim))
+                self.pm.start_virtual_paint_session(channel, current_xlim, current_ylim)
+                
+                # Add first virtual paint stroke
+                brush_size = getattr(self, 'brush_size', 5)
+                self.pm.add_virtual_paint_stroke(x, y, brush_size, self.erase, foreground)
+                
+                # Update display with virtual paint
+                self.pm.update_virtual_paint_display()
         
         elif not self.zoom_mode and event.button == 3:  # Right click (for context menu)
             self.create_context_menu(event)
@@ -2493,92 +2570,6 @@ class ImageViewerWindow(QMainWindow):
             self.selection_start = (event.xdata, event.ydata)
             self.selecting = False  # Will be set to True if the mouse moves while button is held
 
-    def paint_at_position(self, center_x, center_y, erase = False, channel = 2):
-        """Paint pixels within brush radius at given position"""
-        if self.channel_data[channel] is None:
-            return
-
-        if erase:
-            val = 0
-        elif self.machine_window is None:
-            try:
-                val = max(255, self.min_max[4][1])
-            except:
-                val = 255
-        elif self.foreground:
-            val = 1
-        else:
-            val = 2
-        height, width = self.channel_data[channel][self.current_slice].shape
-        radius = self.brush_size // 2
-
-        # Calculate brush area
-        for y in range(max(0, center_y - radius), min(height, center_y + radius + 1)):
-            for x in range(max(0, center_x - radius), min(width, center_x + radius + 1)):
-                # Check if point is within circular brush area
-                if (x - center_x) * 2 + (y - center_y) * 2 <= radius ** 2:
-                    if self.threed and self.threedthresh > 1:
-                        amount = (self.threedthresh - 1) / 2
-                        low = max(0, self.current_slice - amount)
-                        high = min(self.channel_data[channel].shape[0] - 1, self.current_slice + amount)
-                        for i in range(int(low), int(high + 1)):
-                            self.channel_data[channel][i][y, x] = val
-                    else:
-                        self.channel_data[channel][self.current_slice][y, x] = val
-
-    def paint_at_position_vectorized(self, center_x, center_y, erase=False, channel=2, 
-                                   slice_idx=None, brush_size=None, threed=False, 
-                                   threedthresh=1, foreground=True, machine_window=None):
-        """Vectorized paint operation for better performance."""
-        if self.channel_data[channel] is None:
-            return
-        
-        # Use provided parameters or fall back to instance variables
-        slice_idx = slice_idx if slice_idx is not None else self.current_slice
-        brush_size = brush_size if brush_size is not None else getattr(self, 'brush_size', 5)
-        
-        # Determine paint value
-        if erase:
-            val = 0
-        elif machine_window is None:
-            try:
-                val = max(255, self.min_max[4][1])
-            except:
-                val = 255
-        elif foreground:
-            val = 1
-        else:
-            val = 2
-        
-        height, width = self.channel_data[channel][slice_idx].shape
-        radius = brush_size // 2
-        
-        # Calculate affected region bounds
-        y_min = max(0, center_y - radius)
-        y_max = min(height, center_y + radius + 1)
-        x_min = max(0, center_x - radius)
-        x_max = min(width, center_x + radius + 1)
-        
-        if y_min >= y_max or x_min >= x_max:
-            return  # No valid region to paint
-        
-        # Create coordinate grids for the affected region
-        y_coords, x_coords = np.mgrid[y_min:y_max, x_min:x_max]
-        
-        # Calculate distances squared (avoid sqrt for performance)
-        distances_sq = (x_coords - center_x) ** 2 + (y_coords - center_y) ** 2
-        mask = distances_sq <= radius ** 2
-        
-        # Apply paint to affected slices
-        if threed and threedthresh > 1:
-            amount = (threedthresh - 1) / 2
-            low = max(0, int(slice_idx - amount))
-            high = min(self.channel_data[channel].shape[0] - 1, int(slice_idx + amount))
-            
-            for i in range(low, high + 1):
-                self.channel_data[channel][i][y_min:y_max, x_min:x_max][mask] = val
-        else:
-            self.channel_data[channel][slice_idx][y_min:y_max, x_min:x_max][mask] = val
 
     def handle_can(self, x, y):
 
@@ -2648,6 +2639,7 @@ class ImageViewerWindow(QMainWindow):
             return
             
         current_time = time.time()
+        self.rect_time = current_time
         
         if self.selection_start and not self.selecting and not self.pan_mode and not self.brush_mode:
             if (abs(event.xdata - self.selection_start[0]) > 1 or 
@@ -2718,129 +2710,24 @@ class ImageViewerWindow(QMainWindow):
             if event.inaxes != self.ax:
                 return
             
-            # OPTIMIZED: Queue paint operation instead of immediate execution
-            self.queue_paint_operation(event)
+            # Throttle updates like selection rectangle
+            current_time = time.time()
+            if current_time - getattr(self, 'last_paint_update_time', 0) < 0.016:  # ~60fps
+                return
+            self.last_paint_update_time = current_time
             
-            # OPTIMIZED: Schedule display update at controlled frequency
-            if not self.pending_paint_update:
-                self.pending_paint_update = True
-                self.paint_timer.start(16)  # ~60fps max update rate
-
-    def queue_paint_operation(self, event):
-        """Queue a paint operation for background processing."""
-        x, y = int(event.xdata), int(event.ydata)
-
-        if self.pen_button.isChecked():
-            channel = self.active_channel
-        else:
-            channel = 2
-
-        if self.channel_data[channel] is not None:
-            # Prepare paint session if needed
-            if not self.paint_session_active:
-                self.prepare_paint_session(channel)
-            
-            # Create paint operation
-            paint_op = {
-                'type': 'stroke',
-                'x': x,
-                'y': y,
-                'last_pos': getattr(self, 'last_paint_pos', None),
-                'brush_size': self.brush_size,
-                'erase': self.erase,
-                'channel': channel,
-                'slice': self.current_slice,
-                'threed': getattr(self, 'threed', False),
-                'threedthresh': getattr(self, 'threedthresh', 1),
-                'foreground': getattr(self, 'foreground', True),
-                'machine_window': getattr(self, 'machine_window', None)
-            }
+            x, y = int(event.xdata), int(event.ydata)
             
-            # Queue the operation
-            try:
-                self.paint_queue.put_nowait(paint_op)
-            except queue.Full:
-                pass  # Skip if queue is full to avoid blocking
+            # Determine foreground/background for machine window mode
+            foreground = getattr(self, 'foreground', True)
             
-            self.last_paint_pos = (x, y)
-
-    def prepare_paint_session(self, channel):
-        """Prepare optimized background for blitting during paint session."""
-        if self.paint_session_active:
-            return
+            # Add virtual paint stroke with interpolation
+            brush_size = getattr(self, 'brush_size', 5)
+            self.pm.add_virtual_paint_stroke(x, y, brush_size, self.erase, foreground)
             
-        # IMPORTANT: Don't capture background here - let the main display update handle it
-        # We'll capture the background after the proper channel visibility setup
-        self.paint_session_active = True
-
-    def end_paint_session(self):
-        """Clean up after paint session."""
-        self.paint_session_active = False
-        self.last_paint_pos = None
-
-    def paint_worker_loop(self):
-        """Background thread for processing paint operations."""
-        while True:
-            try:
-                paint_op = self.paint_queue.get(timeout=1.0)
-                if paint_op is None:  # Shutdown signal
-                    break
-                
-                with self.paint_lock:
-                    self.execute_paint_operation(paint_op)
-                    
-            except queue.Empty:
-                continue
-
-    def shutdown(self):
-        """Clean shutdown of worker thread."""
-        self.paint_queue.put(None)  # Signal worker to stop
-        if hasattr(self, 'paint_worker'):
-            self.paint_worker.join(timeout=1.0)
-
-    def execute_paint_operation(self, paint_op):
-        """Execute a single paint operation on the data arrays."""
-        if paint_op['type'] == 'stroke':
-            channel = paint_op['channel']
-            x, y = paint_op['x'], paint_op['y']
-            last_pos = paint_op['last_pos']
-            
-            if last_pos is not None:
-                # Paint line from last position to current
-                points = self.get_line_points(last_pos[0], last_pos[1], x, y)
-                for px, py in points:
-                    height, width = self.channel_data[channel][paint_op['slice']].shape
-                    if 0 <= px < width and 0 <= py < height:
-                        self.paint_at_position_vectorized(
-                            px, py, paint_op['erase'], paint_op['channel'],
-                            paint_op['slice'], paint_op['brush_size'],
-                            paint_op['threed'], paint_op['threedthresh'],
-                            paint_op['foreground'], paint_op['machine_window']
-                        )
-            else:
-                # Single point paint
-                height, width = self.channel_data[channel][paint_op['slice']].shape
-                if 0 <= x < width and 0 <= y < height:
-                    self.paint_at_position_vectorized(
-                        x, y, paint_op['erase'], paint_op['channel'],
-                        paint_op['slice'], paint_op['brush_size'],
-                        paint_op['threed'], paint_op['threedthresh'],
-                        paint_op['foreground'], paint_op['machine_window']
-                    )
+            # Update display with virtual paint (super fast)
+            self.pm.update_virtual_paint_display()
 
-    def flush_paint_updates(self):
-        """Update the display with batched paint changes."""
-        self.pending_paint_update = False
-        
-        # Determine which channel to update
-        channel = self.active_channel if hasattr(self, 'pen_button') and self.pen_button.isChecked() else 2
-        
-        # Get current zoom to preserve it
-        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
-        
-        # Update display
-        self.update_display_slice_optimized(channel, preserve_zoom=(current_xlim, current_ylim))
 
     def create_pan_background(self):
         """Create a static background image from currently visible channels with proper rendering"""
@@ -3062,6 +2949,12 @@ class ImageViewerWindow(QMainWindow):
 
     def on_mouse_release(self, event):
         """Handle mouse release events"""
+
+        if self.zoom_mode:
+            rect_condition = (time.time() - self.rect_time) > 0.01 # This is just to prevent non-deliberate rectangle zooming
+        else:
+            rect_condition = True
+
         if self.pan_mode:
             
             self.panning = False
@@ -3069,7 +2962,7 @@ class ImageViewerWindow(QMainWindow):
             self.canvas.setCursor(Qt.CursorShape.OpenHandCursor)
             
         elif event.button == 1:  # Left button release
-            if self.selecting and self.selection_rect is not None:
+            if rect_condition and self.selecting and self.selection_rect is not None:
                 # Get the rectangle bounds
                 x0 = min(self.selection_start[0], event.xdata)
                 y0 = min(self.selection_start[1], event.ydata)
@@ -3077,7 +2970,13 @@ class ImageViewerWindow(QMainWindow):
                 height = abs(event.ydata - self.selection_start[1])
                 shift_pressed = 'shift' in event.modifiers
 
-                if shift_pressed or self.zoom_mode: #Optional targeted zoom
+                if shift_pressed:
+
+                    args = int(x0), int(x0 + width), int(y0), int(y0 + height)
+
+                    self.show_crop_dialog(args)
+
+                elif self.zoom_mode: #Optional targeted zoom
                     
                     self.ax.set_xlim([x0, x0 + width])
                     self.ax.set_ylim([y0 + height, y0])
@@ -3205,32 +3104,23 @@ class ImageViewerWindow(QMainWindow):
 
         #  Handle brush mode cleanup with paint session management
         if self.brush_mode and hasattr(self, 'painting') and self.painting:
+            # Finish current operation
+            self.pm.finish_current_virtual_operation()
+            
+            # Reset last position for next stroke
+            self.last_virtual_pos = None
+            
+            # End this stroke but keep session active for continuous painting
             self.painting = False
 
-            if self.erase:
-                # Restore hidden channels
-                try:
-                    for i in self.restore_channels:
-                        self.channel_visible[i] = True
-                    self.restore_channels = []
-                except:
-                    pass
-                
-                self.end_paint_session()
-                
-                # OPTIMIZED: Stop timer and process any pending paint operations
-                if hasattr(self, 'paint_timer'):
-                    self.paint_timer.stop()
-                if hasattr(self, 'pending_paint_update') and self.pending_paint_update:
-                    self.flush_paint_updates()
-
-                self.static_background = None
-
-                current_xlim = self.ax.get_xlim() if hasattr(self, 'ax') and self.ax.get_xlim() != (0, 1) else None
+            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
 
-                current_ylim = self.ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
+            self.update_display(preserve_zoom=(current_xlim, current_ylim), continue_paint = True)
 
-                self.update_display(preserve_zoom = (current_xlim, current_ylim))
+            if self.machine_window is not None:
+                if self.machine_window.segmentation_worker is not None:
+                    self.machine_window.segmentation_worker.resume()
 
             
 
@@ -3539,7 +3429,7 @@ class ImageViewerWindow(QMainWindow):
         ripley_action.triggered.connect(self.show_ripley_dialog)
         heatmap_action = stats_menu.addAction("Community Cluster Heatmap")
         heatmap_action.triggered.connect(self.show_heatmap_dialog)
-        nearneigh_action = stats_menu.addAction("Average Nearest Neighbors")
+        nearneigh_action = stats_menu.addAction("Average Nearest Neighbors (With Clustering Heatmaps)")
         nearneigh_action.triggered.connect(self.show_nearneigh_dialog)
         vol_action = stats_menu.addAction("Calculate Volumes")
         vol_action.triggered.connect(self.volumes)
@@ -3558,6 +3448,8 @@ class ImageViewerWindow(QMainWindow):
         community_code_action.triggered.connect(lambda: self.show_code_dialog(sort = 'Community'))
         id_code_action = overlay_menu.addAction("Code Identities")
         id_code_action.triggered.connect(lambda: self.show_code_dialog(sort = 'Identity'))
+        umap_action = overlay_menu.addAction("Centroid UMAP")
+        umap_action.triggered.connect(self.handle_umap)
 
         rand_menu = analysis_menu.addMenu("Randomize")
         random_action = rand_menu.addAction("Generate Equivalent Random Network")
@@ -3574,6 +3466,10 @@ class ImageViewerWindow(QMainWindow):
         calc_all_action.triggered.connect(self.show_calc_all_dialog)
         calc_prox_action = calculate_menu.addAction("Calculate Proximity Network (connect nodes by distance)")
         calc_prox_action.triggered.connect(self.show_calc_prox_dialog)
+        calc_branch_action = calculate_menu.addAction("Calculate Branchpoint Network (Connect Branchpoints of Edge Image - Good for Nerves/Vessels)")
+        calc_branch_action.triggered.connect(self.handle_calc_branch)
+        calc_branchprox_action = calculate_menu.addAction("Calculate Branch Adjacency Network (Of Edges)")
+        calc_branchprox_action.triggered.connect(self.handle_branchprox_calc)
         centroid_action = calculate_menu.addAction("Calculate Centroids (Active Image)")
         centroid_action.triggered.connect(self.show_centroid_dialog)
 
@@ -3597,13 +3493,17 @@ class ImageViewerWindow(QMainWindow):
         mask_action = image_menu.addAction("Mask Channel")
         mask_action.triggered.connect(self.show_mask_dialog)
         crop_action = image_menu.addAction("Crop Channels")
-        crop_action.triggered.connect(self.show_crop_dialog)
+        crop_action.triggered.connect(lambda: self.show_crop_dialog(args = None))
         type_action = image_menu.addAction("Channel dtype")
         type_action.triggered.connect(self.show_type_dialog)
         skeletonize_action = image_menu.addAction("Skeletonize")
         skeletonize_action.triggered.connect(self.show_skeletonize_dialog)
-        watershed_action = image_menu.addAction("Watershed")
+        dt_action = image_menu.addAction("Distance Transform (For binary images)")
+        dt_action.triggered.connect(self.show_dt_dialog)
+        watershed_action = image_menu.addAction("Binary Watershed")
         watershed_action.triggered.connect(self.show_watershed_dialog)
+        gray_water_action = image_menu.addAction("Gray Watershed")
+        gray_water_action.triggered.connect(self.show_gray_water_dialog)
         invert_action = image_menu.addAction("Invert")
         invert_action.triggered.connect(self.show_invert_dialog)
         z_proj_action = image_menu.addAction("Z Project")
@@ -3615,7 +3515,7 @@ class ImageViewerWindow(QMainWindow):
         gennodes_action = generate_menu.addAction("Generate Nodes (From 'Edge' Vertices)")
         gennodes_action.triggered.connect(self.show_gennodes_dialog)
         branch_action = generate_menu.addAction("Label Branches")
-        branch_action.triggered.connect(self.show_branch_dialog)
+        branch_action.triggered.connect(lambda: self.show_branch_dialog())
         genvor_action = generate_menu.addAction("Generate Voronoi Diagram (From Node Centroids) - goes in Overlay2")
         genvor_action.triggered.connect(self.voronoi)
 
@@ -3654,6 +3554,43 @@ class ImageViewerWindow(QMainWindow):
         help_button = menubar.addAction("Help")
         help_button.triggered.connect(self.help_me)
 
+        cam_button = QPushButton("📷")
+        cam_button.setFixedSize(40, 40)
+        cam_button.setStyleSheet("font-size: 24px;")  # Makes emoji larger
+        cam_button.clicked.connect(self.snap)
+        menubar.setCornerWidget(cam_button, Qt.Corner.TopRightCorner)
+
+    def snap(self):
+
+        try:
+
+            for thing in self.channel_data:
+                if thing is not None:
+                    data = True
+            if not data:
+                return
+
+            snap = self.create_composite_for_pan()
+
+            filename, _ = QFileDialog.getSaveFileName(
+                self,
+                f"Save Image As",
+                "",  # Default directory
+                "TIFF Files (*.tif *.tiff);;All Files (*)"  # File type filter
+            )
+            
+            if filename:  # Only proceed if user didn't cancel
+                # If user didn't type an extension, add .tif
+                if not filename.endswith(('.tif', '.tiff')):
+                    filename += '.tif'
+
+            import tifffile
+            tifffile.imwrite(filename, snap)
+
+        except:
+            pass
+
+
     def open_cellpose(self):
 
         try:
@@ -3919,6 +3856,10 @@ class ImageViewerWindow(QMainWindow):
         dialog = MergeNodeIdDialog(self)
         dialog.exec()
 
+    def show_gray_water_dialog(self):
+        """Show the gray watershed parameter dialog."""
+        dialog = GrayWaterDialog(self)
+        dialog.exec()
 
     def show_watershed_dialog(self):
         """Show the watershed parameter dialog."""
@@ -3950,6 +3891,110 @@ class ImageViewerWindow(QMainWindow):
         dialog = ProxDialog(self)
         dialog.exec()
 
+    def table_load_attrs(self):
+
+        # Display network_lists in the network table
+        try:
+            if hasattr(my_network, 'network_lists'):
+                model = PandasModel(my_network.network_lists)
+                self.network_table.setModel(model)
+                # Adjust column widths to content
+                for column in range(model.columnCount(None)):
+                    self.network_table.resizeColumnToContents(column)
+        except Exception as e:
+            print(f"Error loading network_lists: {e}")
+
+        #Display the other things if they exist
+        try:
+
+            if hasattr(my_network, 'node_identities') and my_network.node_identities is not None:
+                try:
+                    self.format_for_upperright_table(my_network.node_identities, 'NodeID', 'Identity', 'Node Identities')
+                except Exception as e:
+                    print(f"Error loading node identity table: {e}")
+
+            if hasattr(my_network, 'node_centroids') and my_network.node_centroids is not None:
+                try:
+                    self.format_for_upperright_table(my_network.node_centroids, 'NodeID', ['Z', 'Y', 'X'], 'Node Centroids')
+                except Exception as e:
+                    print(f"Error loading node centroid table: {e}")
+
+
+            if hasattr(my_network, 'edge_centroids') and my_network.edge_centroids is not None:
+                try:
+                    self.format_for_upperright_table(my_network.edge_centroids, 'EdgeID', ['Z', 'Y', 'X'], 'Edge Centroids')
+                except Exception as e:
+                    print(f"Error loading edge centroid table: {e}")
+
+
+        except Exception as e:
+            print(f"An error has occured: {e}")
+
+    def confirm_calcbranch_dialog(self, message):
+        """Shows a dialog asking user to confirm if they want to proceed below"""
+        msg = QMessageBox()
+        msg.setIcon(QMessageBox.Icon.Question)
+        msg.setText("Alert")
+        msg.setInformativeText(message)
+        msg.setWindowTitle("Proceed?")
+        msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
+        return msg.exec() == QMessageBox.StandardButton.Yes
+
+    def handle_calc_branch(self):
+
+        try:
+
+            if self.channel_data[0] is not None or self.channel_data[3] is not None:
+                if not self.confirm_calcbranch_dialog("Use of this feature will require additional use of the Nodes and Overlay 2 channels. Please save any data and return, or proceed if you do not need those channels' data"):
+                    return
+
+            my_network.id_overlay = my_network.edges.copy()
+
+            self.show_gennodes_dialog()
+
+            my_network.edges = (my_network.nodes == 0) * my_network.edges
+
+            my_network.calculate_all(my_network.nodes, my_network.edges, xy_scale = my_network.xy_scale, z_scale = my_network.z_scale, search = None, diledge = None, inners = False, hash_inners = False, remove_trunk = 0, ignore_search_region = True, other_nodes = None, label_nodes = True, directory = None, GPU = False, fast_dil = False, skeletonize = False, GPU_downsample = None)
+
+            self.load_channel(1, my_network.edges, data = True)
+            self.load_channel(0, my_network.nodes, data = True)
+            self.load_channel(3, my_network.id_overlay, data = True)
+
+            self.table_load_attrs()
+
+        except Exception as e:
+
+            try:
+                my_network.edges = my_network.id_overlay
+                my_network.id_overlay = None
+            except:
+                pass
+
+            print(f"Error calculating branchpoint network: {e}")
+
+    def handle_branchprox_calc(self):
+
+        try:
+
+            if self.channel_data[0] is not None:
+                if not self.confirm_calcbranch_dialog("Use of this feature will require additional use of the Nodes and Overlay 2 channels. Please save any data and return, or proceed if you do not need those channels' data"):
+                    return
+
+            self.show_branch_dialog(called = True)
+
+            self.load_channel(0, my_network.edges, data = True)
+
+            self.delete_channel(1, False)
+
+            my_network.morph_proximity(search = [3,3], fastdil = True)
+
+            self.table_load_attrs()
+
+        except Exception as e:
+
+            print(f"Error calculating network: {e}")
+
+
     def show_centroid_dialog(self):
         """show the centroid dialog"""
         dialog = CentroidDialog(self)
@@ -3994,9 +4039,9 @@ class ImageViewerWindow(QMainWindow):
         dialog = MaskDialog(self)
         dialog.exec()
 
-    def show_crop_dialog(self):
+    def show_crop_dialog(self, args = None):
         """Show the crop dialog"""
-        dialog = CropDialog(self)
+        dialog = CropDialog(self, args = args)
         dialog.exec()
 
     def show_type_dialog(self):
@@ -4012,6 +4057,11 @@ class ImageViewerWindow(QMainWindow):
         dialog = SkeletonizeDialog(self)
         dialog.exec()
 
+    def show_dt_dialog(self):
+        """show the dt dialog"""
+        dialog = DistanceDialog(self)
+        dialog.exec()
+
     def show_centroid_node_dialog(self):
         """show the centroid node dialog"""
         dialog = CentroidNodeDialog(self)
@@ -4023,9 +4073,9 @@ class ImageViewerWindow(QMainWindow):
         gennodes = GenNodesDialog(self, down_factor = down_factor, called = called)
         gennodes.exec()
 
-    def show_branch_dialog(self):
+    def show_branch_dialog(self, called = False):
         """Show the branch label dialog"""
-        dialog = BranchDialog(self)
+        dialog = BranchDialog(self, called = called)
         dialog.exec()
 
     def voronoi(self):
@@ -4178,6 +4228,7 @@ class ImageViewerWindow(QMainWindow):
                     if sort == 'Node Identities':
                         my_network.load_node_identities(file_path = filename)
 
+                        """
                         first_value = list(my_network.node_identities.values())[0]  # Check that there are not multiple IDs
                         if isinstance(first_value, (list, tuple)):
                             trump_value, ok = QInputDialog.getText(
@@ -4199,7 +4250,7 @@ class ImageViewerWindow(QMainWindow):
                         else:
                             trump_value = None
                             my_network.node_identities = uncork(my_network.node_identities, trump_value)
-
+                        """
 
                         if hasattr(my_network, 'node_identities') and my_network.node_identities is not None:
                             try:
@@ -4316,6 +4367,7 @@ class ImageViewerWindow(QMainWindow):
                     )
 
             self.last_load = directory
+            
 
             if directory != "":
 
@@ -4354,7 +4406,7 @@ class ImageViewerWindow(QMainWindow):
                 # Display network_lists in the network table
                 # Create empty DataFrame for network table if network_lists is None
                 if not hasattr(my_network, 'network_lists') or my_network.network_lists is None:
-                    empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+                    empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
                     model = PandasModel(empty_df)
                     self.network_table.setModel(model)
                 else:
@@ -4439,7 +4491,7 @@ class ImageViewerWindow(QMainWindow):
         """Method to close Excelotron"""
         self.excel_manager.close()
     
-    def handle_excel_data(self, data_dict, property_name):
+    def handle_excel_data(self, data_dict, property_name, add):
         """Handle data received from Excelotron"""
         print(f"Received data for property: {property_name}")
         print(f"Data keys: {list(data_dict.keys())}")
@@ -4448,12 +4500,19 @@ class ImageViewerWindow(QMainWindow):
 
             try:
 
+                if not add or my_network.node_centroids is None:
+                    centroids = {}
+                    max_val = 0
+                else:
+                    centroids = my_network.node_centroids
+                    max_val = max(list(my_network.node_centroids.keys()))
+
                 ys = data_dict['Y']
                 xs = data_dict['X']
                 if 'Numerical IDs' in data_dict:
                     nodes = data_dict['Numerical IDs']
                 else:
-                    nodes = np.arange(1, len(ys) + 1)
+                    nodes = np.arange(max_val + 1, max_val + len(ys) + 1)
 
 
                 if 'Z' in data_dict:
@@ -4461,8 +4520,6 @@ class ImageViewerWindow(QMainWindow):
                 else:
                     zs = np.zeros(len(ys))
 
-                centroids = {}
-
                 for i in range(len(nodes)):
 
                     centroids[nodes[i]] = [int(zs[i]), int(ys[i]), int(xs[i])]
@@ -4480,15 +4537,26 @@ class ImageViewerWindow(QMainWindow):
 
             try:
 
+                if not add or my_network.node_identities is None:
+                    identities = {}
+                    max_val = 0
+                else:
+                    identities = my_network.node_identities
+                    if my_network.node_centroids is not None:
+                        max_val = max(list(my_network.node_centroids.keys()))
+                    else:
+                        max_val = max(list(my_network.node_identities.keys()))
+
                 idens = data_dict['Identity Column']
 
                 if 'Numerical IDs' in data_dict:
                     nodes = data_dict['Numerical IDs']
-                else:
-                    nodes = np.arange(1, len(idens) + 1)
-
-                identities = {}
+                    if add:
+                        for i, node in enumerate(nodes):
+                            nodes[i] = node + max_val
 
+                else:
+                    nodes = np.arange(max_val + 1, max_val + len(data_dict['Identity Column']) + 1)
 
                 for i in range(len(nodes)):
 
@@ -4507,14 +4575,20 @@ class ImageViewerWindow(QMainWindow):
 
             try:
 
+                if not add or my_network.communities is None:
+                    communities = {}
+                    max_val = 0
+                else:
+                    communities = my_network.communities
+                    max_val = max(list(my_network.communities.keys()))
+
+
                 coms = data_dict['Community Identifier']
 
                 if 'Numerical IDs' in data_dict:
                     nodes = data_dict['Numerical IDs']
                 else:
-                    nodes = np.arange(1, len(coms) + 1)
-
-                communities = {}
+                    nodes = np.arange(max_val + 1, max_val + len(data_dict['Community Identifier']) + 1)
 
                 for i in range(len(nodes)):
 
@@ -4556,6 +4630,7 @@ class ImageViewerWindow(QMainWindow):
             - 'mean': averages across color channels
             - 'max': takes maximum value across color channels
             - 'min': takes minimum value across color channels
+            - 'weight': takes weighted channel averages
         
         Returns:
         --------
@@ -4570,7 +4645,7 @@ class ImageViewerWindow(QMainWindow):
         if array.ndim != 4:
             raise ValueError(f"Expected 4D array, got {array.ndim}D array")
         
-        if method not in ['first', 'mean', 'max', 'min']:
+        if method not in ['first', 'mean', 'max', 'min', 'weight']:
             raise ValueError(f"Unknown method: {method}")
         
         if method == 'first':
@@ -4579,6 +4654,9 @@ class ImageViewerWindow(QMainWindow):
             return np.mean(array, axis=-1)
         elif method == 'max':
             return np.max(array, axis=-1)
+        elif method == 'weight':
+            # Apply the luminosity formula
+            return (0.2989 * array[:,:,:,0] + 0.5870 * array[:,:,:,1] + 0.1140 * array[:,:,:,2])
         else:  # min
             return np.min(array, axis=-1)
 
@@ -4602,7 +4680,7 @@ class ImageViewerWindow(QMainWindow):
         msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
         return msg.exec() == QMessageBox.StandardButton.Yes
 
-    def load_channel(self, channel_index, channel_data=None, data=False, assign_shape = True, preserve_zoom = None, end_paint = False, begin_paint = False):
+    def load_channel(self, channel_index, channel_data=None, data=False, assign_shape = True, preserve_zoom = None, end_paint = False, begin_paint = False, color = False):
         """Load a channel and enable active channel selection if needed."""
 
         try:
@@ -4627,7 +4705,6 @@ class ImageViewerWindow(QMainWindow):
                         import tifffile
                         self.channel_data[channel_index] = tifffile.imread(filename)
 
-                        
                     elif file_extension == 'nii':
                         import nibabel as nib
                         nii_img = nib.load(filename)
@@ -4657,9 +4734,20 @@ class ImageViewerWindow(QMainWindow):
                     self.channel_buttons[channel_index].setEnabled(False)
                     self.delete_buttons[channel_index].setEnabled(False) 
 
+            try:
+                #if len(self.channel_data[channel_index].shape) == 4:
+                if 1 in self.channel_data[channel_index].shape:
+                    print("Removing singleton dimension (I am assuming this is a channel dimension?)")
+                    self.channel_data[channel_index] = np.squeeze(self.channel_data[channel_index])
+            except:
+                pass
+
             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)
 
+            if self.channel_data[channel_index].dtype == np.bool_: #Promote boolean arrays if they somehow get loaded
+                self.channel_data[channel_index] = self.channel_data[channel_index].astype(np.uint8)
+
             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
@@ -4668,12 +4756,13 @@ class ImageViewerWindow(QMainWindow):
                             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
+
+            if not color:
+                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], 'weight')
+                except:
+                    pass
 
             reset_resize = False
 
@@ -4748,8 +4837,13 @@ class ImageViewerWindow(QMainWindow):
             if all(not btn.isEnabled() for btn in self.channel_buttons[:channel_index]):
                 self.set_active_channel(channel_index)
 
-            if not self.channel_buttons[channel_index].isChecked():
-                self.channel_buttons[channel_index].click()
+            if self.chan_load:
+                if not self.channel_buttons[channel_index].isChecked():
+                    self.channel_buttons[channel_index].click()
+            else:
+                if self.channel_buttons[channel_index].isChecked():
+                    self.channel_buttons[channel_index].click()
+
             self.min_max[channel_index][0] = np.min(self.channel_data[channel_index])
             self.min_max[channel_index][1] = np.max(self.channel_data[channel_index])
             self.volume_dict[channel_index] = None #reset volumes
@@ -4849,7 +4943,7 @@ class ImageViewerWindow(QMainWindow):
             my_network.communities = None
 
             # Create empty DataFrame
-            empty_df = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+            empty_df = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
             
             # Clear network table
             self.network_table.setModel(PandasModel(empty_df))
@@ -5018,25 +5112,29 @@ class ImageViewerWindow(QMainWindow):
 
 
     
-    def update_display(self, preserve_zoom=None, dims = None, called = False, reset_resize = False):
+    def update_display(self, preserve_zoom=None, dims = None, called = False, reset_resize = False, continue_paint = False, skip_paint_reinit = False):
         """Update the display with currently visible channels and highlight overlay."""
-
         try:
-
             self.figure.clear()
-
             if self.pan_background_image is not None:
                 # Restore previously visible channels
                 self.channel_visible = self.pre_pan_channel_state.copy()
                 self.is_pan_preview = False
                 self.pan_background_image = None
-
                 if self.machine_window is not None:
                     if self.machine_window.segmentation_worker is not None:
                         self.machine_window.segmentation_worker.resume()
-
             if self.static_background is not None:
-
+                # NEW: Convert virtual strokes to real data before cleanup
+                if (hasattr(self, 'virtual_draw_operations') and self.virtual_draw_operations) or \
+                   (hasattr(self, 'virtual_erase_operations') and self.virtual_erase_operations) or \
+                   (hasattr(self, 'current_operation') and self.current_operation):
+                    # Finish current operation first
+                    if hasattr(self, 'current_operation') and self.current_operation:
+                        self.pm.finish_current_virtual_operation()
+                    # Now convert to real data
+                    self.pm.convert_virtual_strokes_to_data()
+                
                 # Restore hidden channels
                 try:
                     for i in self.restore_channels:
@@ -5044,27 +5142,17 @@ class ImageViewerWindow(QMainWindow):
                     self.restore_channels = []
                 except:
                     pass
-                
-                self.end_paint_session()
-                
-                # OPTIMIZED: Stop timer and process any pending paint operations
-                if hasattr(self, 'paint_timer'):
-                    self.paint_timer.stop()
-                if hasattr(self, 'pending_paint_update') and self.pending_paint_update:
-                    self.flush_paint_updates()
-
-                self.static_background = None
-
-                if self.machine_window is None:
-
-                    try:
+                if not continue_paint:
+                    self.static_background = None
 
-                        self.channel_data[4][self.current_slice, :, :] = n3d.overlay_arrays_simple(self.channel_data[self.temp_chan][self.current_slice, :, :], self.channel_data[4][self.current_slice, :, :])
-                        self.load_channel(self.temp_chan, self.channel_data[4], data = True, end_paint = True)
-                        self.channel_data[4] = None
-                        self.channel_visible[4] = False
-                    except:
-                        pass
+                    if self.machine_window is None:
+                        try:
+                            self.channel_data[4][self.current_slice, :, :] = n3d.overlay_arrays_simple(self.channel_data[self.temp_chan][self.current_slice, :, :], self.channel_data[4][self.current_slice, :, :])
+                            self.load_channel(self.temp_chan, self.channel_data[4], data = True, end_paint = True)
+                            self.channel_data[4] = None
+                            self.channel_visible[4] = False
+                        except:
+                            pass
 
             # Get active channels and their dimensions
             active_channels = [i for i in range(4) if self.channel_data[i] is not None]
@@ -5295,49 +5383,22 @@ class ImageViewerWindow(QMainWindow):
             
             self.canvas.draw()
 
+            if self.brush_mode and not skip_paint_reinit:
+                # Get current zoom to preserve it
+                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 self.pen_button.isChecked():
+                    channel = self.active_channel
+                else:
+                    channel = 2
+
+                self.pm.initiate_paint_session(channel, current_xlim, current_ylim)
+
         except:
             import traceback
             print(traceback.format_exc())
 
-    def update_display_slice_optimized(self, channel, preserve_zoom=None):
-        """Ultra minimal update that only changes the paint channel's data - OPTIMIZED VERSION"""
-        if not self.channel_visible[channel]:
-            return
-            
-        if preserve_zoom:
-            current_xlim, current_ylim = preserve_zoom
-            if current_xlim is not None and current_ylim is not None:
-                self.ax.set_xlim(current_xlim)
-                self.ax.set_ylim(current_ylim)
-        
-        # Find the existing image for channel (paint channel)
-        channel_image = None
-        for img in self.ax.images:
-            if img.cmap.name == f'custom_{channel}':
-                channel_image = img
-                break
-                
-        if channel_image is not None:
-            # Update the data of the existing image with thread safety
-            with self.paint_lock:
-                channel_image.set_array(self.channel_data[channel][self.current_slice])
-            
-            # Restore the static background (all other channels) at current zoom level
-            # This is the key - use static_background from update_display, not paint_background
-            if hasattr(self, 'static_background') and self.static_background is not None:
-                self.canvas.restore_region(self.static_background)
-                # Draw just our paint channel
-                self.ax.draw_artist(channel_image)
-                # Blit everything
-                self.canvas.blit(self.ax.bbox)
-                self.canvas.flush_events()
-            else:
-                # Fallback to full draw if no static background
-                self.canvas.draw()
-        else:
-            # Fallback if channel image not found
-            self.canvas.draw()
-
     def get_channel_image(self, channel):
         """Find the matplotlib image object for a specific channel."""
         if not hasattr(self.ax, 'images'):
@@ -5467,6 +5528,13 @@ class ImageViewerWindow(QMainWindow):
         dialog = CodeDialog(self, sort = sort)
         dialog.exec()
 
+    def handle_umap(self):
+
+        if my_network.node_centroids is None:
+            self.show_centroid_dialog()
+
+        my_network.centroid_umap()
+
     def closeEvent(self, event):
         """Override closeEvent to close all windows when main window closes"""
         
@@ -6009,9 +6077,9 @@ class PandasModel(QAbstractTableModel):
         if data is None:
             # Create an empty DataFrame with default columns
             import pandas as pd
-            data = pd.DataFrame(columns=['Node 1A', 'Node 1B', 'Edge 1C'])
+            data = pd.DataFrame(columns=['Node A', 'Node B', 'Edge C'])
         elif type(data) == list:
-            data = self.lists_to_dataframe(data[0], data[1], data[2], column_names=['Node 1A', 'Node 1B', 'Edge 1C'])
+            data = self.lists_to_dataframe(data[0], data[1], data[2], column_names=['Node A', 'Node B', 'Edge C'])
         self._data = data
         self.bold_cells = set()
         self.highlighted_cells = set()
@@ -6752,7 +6820,7 @@ class MergeNodeIdDialog(QDialog):
 
             self.parent().format_for_upperright_table(my_network.node_identities, 'NodeID', 'Identity')
 
-            QMessageBox.critical(
+            QMessageBox.information(
                 self,
                 "Success",
                 "Node Identities Merged. New IDs represent presence of corresponding img foreground with +, absence with -. Please save your new identities as csv, then use File -> Load -> Load From Excel Helper to bulk search and rename desired combinations. (Press Help [above] for more info)"
@@ -6987,24 +7055,29 @@ class ColorOverlayDialog(QDialog):
 
     def coloroverlay(self):
 
-        down_factor = float(self.down_factor.text()) if self.down_factor.text().strip() else None
+        try:
 
-        if self.parent().active_channel == 0:
-            mode = 0
-            self.sort = 'Node'
-        else:
-            mode = 1
-            self.sort = 'Edge'
+            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)
+            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().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.parent().load_channel(3, channel_data = result, data = True)
 
-        self.accept()
+            self.accept()
+
+        except:
+            pass
 
 
 class ShuffleDialog(QDialog):
@@ -7211,6 +7284,8 @@ class PartitionDialog(QDialog):
 
     def partition(self):
 
+        self.parent().prev_coms = None
+
         accepted_mode = self.mode_selector.currentIndex()
         weighted = self.weighted.isChecked()
         dostats = self.stats.isChecked()
@@ -7345,7 +7420,7 @@ class ComNeighborDialog(QDialog):
         # weighted checkbox (default True)
         self.proportional = QPushButton("Robust")
         self.proportional.setCheckable(True)
-        self.proportional.setChecked(False)
+        self.proportional.setChecked(True)
         layout.addRow("Return Node Type Distribution Robust Heatmaps (ie, will give two more heatmaps that are not beholden to the total number of nodes of each type, representing which structures are overrepresented in a network):", self.proportional)
 
         self.mode = QComboBox()
@@ -7393,7 +7468,7 @@ class ComNeighborDialog(QDialog):
                 self.parent().format_for_upperright_table(matrix, 'NeighborhoodID', id_set, title = f'Neighborhood Heatmap {i + 1}')
 
 
-            self.parent().format_for_upperright_table(len_dict, 'NeighborhoodID', 'Proportion of Total Nodes', title = 'Neighborhood Counts')
+            self.parent().format_for_upperright_table(len_dict, 'NeighborhoodID', ['Number of Communities', 'Proportion of Total Nodes'], title = 'Neighborhood Counts')
             self.parent().format_for_upperright_table(my_network.communities, 'NodeID', 'NeighborhoodID', title = 'Neighborhood Partition')
 
             print("Neighborhoods have been assigned to communities based on similarity")
@@ -7435,6 +7510,8 @@ class ComCellDialog(QDialog):
 
         try:
 
+            self.parent().prev_coms = None
+
             size = float(self.size.text()) if self.size.text().strip() else None
             xy_scale = float(self.xy_scale.text()) if self.xy_scale.text().strip() else 1
             z_scale = float(self.z_scale.text()) if self.z_scale.text().strip() else 1
@@ -7596,6 +7673,16 @@ class NearNeighDialog(QDialog):
         heatmap_layout.addRow("Overlay:", self.numpy)
         
         main_layout.addWidget(heatmap_group)
+
+        quant_group = QGroupBox("Quantifiable Overlay")
+        quant_layout = QFormLayout(quant_group)
+
+        self.quant = QPushButton("Return quantifiable overlay? (Labels nodes by distance, good with intensity-thresholding to isolate targets. Requires labeled nodes image.)")
+        self.quant.setCheckable(True)
+        self.quant.setChecked(False)
+        quant_layout.addRow("Overlay:", self.quant)
+
+        main_layout.addWidget(quant_group)
         
         # Get Distribution group box
         distribution_group = QGroupBox("Get Distribution")
@@ -7643,6 +7730,7 @@ class NearNeighDialog(QDialog):
             threed = self.threed.isChecked()
             numpy = self.numpy.isChecked()
             num = int(self.num.text()) if self.num.text().strip() else 1
+            quant = self.quant.isChecked()
 
             if root is not None and targ is not None:
                 title = f"Nearest {num} Neighbor(s) Distance of {targ} from {root}"
@@ -7659,11 +7747,14 @@ class NearNeighDialog(QDialog):
                     return
 
             if not numpy:
-                avg, output = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num, heatmap = heatmap, threed = threed)
+                avg, output, quant_overlay = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num, heatmap = heatmap, threed = threed, quant = quant)
             else:
-                avg, output, overlay = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num, heatmap = heatmap, threed = threed, numpy = True)
+                avg, output, overlay, quant_overlay = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num, heatmap = heatmap, threed = threed, numpy = True, quant = quant)
                 self.parent().load_channel(3, overlay, data = True)
 
+            if quant_overlay is not None:
+                self.parent().load_channel(2, quant_overlay, data = True)
+            
             self.parent().format_for_upperright_table([avg], metric = f'Avg {title}', title = f'Avg {title}')
             self.parent().format_for_upperright_table(output, header2, header, title = title)
 
@@ -7691,7 +7782,7 @@ class NearNeighDialog(QDialog):
 
                 for targ in available:
 
-                    avg, _ = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num)
+                    avg, _, _ = my_network.nearest_neighbors_avg(root, targ, my_network.xy_scale, my_network.z_scale, num = num)
 
                     output_dict[f"{root} vs {targ}"] = avg
 
@@ -7781,52 +7872,86 @@ class NeighborIdentityDialog(QDialog):
 
 
 class RipleyDialog(QDialog):
-
     def __init__(self, parent=None):
-
         super().__init__(parent)
         self.setWindowTitle(f"Find Ripley's H Function From Centroids")
         self.setModal(True)
-
-        layout = QFormLayout(self)
-
+        
+        # Main layout
+        main_layout = QVBoxLayout(self)
+        
+        # Node Parameters Group (only if node_identities exist)
         if my_network.node_identities is not None:
+            node_group = QGroupBox("Node Parameters")
+            node_layout = QFormLayout(node_group)
+            
             self.root = QComboBox()
             self.root.addItems(list(set(my_network.node_identities.values())))  
             self.root.setCurrentIndex(0)
-            layout.addRow("Root Identity to Search for Neighbors", self.root)
-        else:
-            self.root = None
-
-        if my_network.node_identities is not None:
+            node_layout.addRow("Root Identity to Search for Neighbors:", self.root)
+            
             self.targ = QComboBox()
             self.targ.addItems(list(set(my_network.node_identities.values())))  
             self.targ.setCurrentIndex(0)
-            layout.addRow("Targ Identity to be Searched For", self.targ)
+            node_layout.addRow("Target Identity to be Searched For:", self.targ)
+            
+            main_layout.addWidget(node_group)
         else:
+            self.root = None
             self.targ = None
-
+        
+        # Search Parameters Group
+        search_group = QGroupBox("Search Parameters")
+        search_layout = QFormLayout(search_group)
+        
         self.distance = QLineEdit("5")
-        layout.addRow("Bucket Distance for Searching For Clusters (automatically scaled by xy and z scales):", self.distance)
-
-
+        search_layout.addRow("1. Bucket Distance for Searching For Clusters\n(automatically scaled by xy and z scales):", self.distance)
+        
         self.proportion = QLineEdit("0.5")
-        layout.addRow("Proportion of image to search? (0-1, high vals increase border artifacts): ", self.proportion)
-
+        search_layout.addRow("2. Proportion of image to search?\n(0-1, high vals increase border artifacts):", self.proportion)
+        
+        main_layout.addWidget(search_group)
+        
+        # Border Safety Group
+        border_group = QGroupBox("Border Safety")
+        border_layout = QFormLayout(border_group)
+        
+        self.ignore = QPushButton("Ignore Border Roots")
+        self.ignore.setCheckable(True)
+        self.ignore.setChecked(True)
+        border_layout.addRow("3. Exclude Root Nodes Near Borders?:", self.ignore)
+        
+        self.factor = QLineEdit("0.5")
+        border_layout.addRow("4. (If param 3): Proportion of most internal nodes to use? (0 < n < 1) (Higher = more internal)?:", self.factor)
+        
+        self.mode = QComboBox()
+        self.mode.addItems(["Boundaries of Entire Image", "Boundaries of Edge Image Mask", 
+                           "Boundaries of Overlay1 Mask", "Boundaries of Overlay2 Mask"])
+        self.mode.setCurrentIndex(0)
+        border_layout.addRow("5. (If param 3): Define Boundaries How?:", self.mode)
+        
+        self.safe = QPushButton("Ignore Border Radii")
+        self.safe.setCheckable(True)
+        self.safe.setChecked(True)
+        border_layout.addRow("6. (If param 3): Keep search radii within border (overrides Param 2, also assigns volume to that of mask)?:", self.safe)
+        
+        main_layout.addWidget(border_group)
+        
+        # Experimental Border Safety Group
+        experimental_group = QGroupBox("Aggressive Border Safety (Creates duplicate centroids reflected across the image border - if you really need to search there for whatever reason - Not meant to be used if confining search to a masked object)")
+        experimental_layout = QFormLayout(experimental_group)
+        
         self.edgecorrect = QPushButton("Border Correction")
         self.edgecorrect.setCheckable(True)
         self.edgecorrect.setChecked(False)
-        layout.addRow("Use Border Correction (Extrapolate for points beyond the border):", self.edgecorrect)
-
-        self.ignore = QPushButton("Ignore Border Roots")
-        self.ignore.setCheckable(True)
-        self.ignore.setChecked(False)
-        layout.addRow("Exclude Root Nodes Near Borders?:", self.ignore)
-
+        experimental_layout.addRow("7. Use Border Correction\n(Extrapolate for points beyond the border):", self.edgecorrect)
+        
+        main_layout.addWidget(experimental_group)
+        
         # Add Run button
         run_button = QPushButton("Get Ripley's H")
         run_button.clicked.connect(self.ripley)
-        layout.addWidget(run_button)
+        main_layout.addWidget(run_button)
 
     def ripley(self):
 
@@ -7856,6 +7981,16 @@ class RipleyDialog(QDialog):
             except:
                 proportion = 0.5
 
+            try:
+                factor = abs(float(self.factor.text()))
+
+            except:
+                factor = 0.25
+
+            if factor > 1 or factor <= 0:
+                print("Utilizing factor = 0.25")
+                factor = 0.25
+
             if proportion > 1 or proportion <= 0:
                 print("Utilizing proportion = 0.5")
                 proportion = 0.5
@@ -7865,6 +8000,13 @@ class RipleyDialog(QDialog):
 
             ignore = self.ignore.isChecked()
 
+            safe = self.safe.isChecked()
+
+            mode = self.mode.currentIndex()
+
+            if mode == 0:
+                factor = factor/2 #The logic treats this as distance to border later, only if mode is 0, but its supposed to represent proportion internal.
+
             if my_network.nodes is not None:
 
                 if my_network.nodes.shape[0] == 1:
@@ -7874,7 +8016,7 @@ class RipleyDialog(QDialog):
             else:
                 bounds = None
 
-            r_vals, k_vals, h_vals = my_network.get_ripley(root, targ, distance, edgecorrect, bounds, ignore, proportion)
+            r_vals, k_vals, h_vals = my_network.get_ripley(root, targ, distance, edgecorrect, bounds, ignore, proportion, mode, safe, factor)
             
             k_dict = dict(zip(r_vals, k_vals))
             h_dict = dict(zip(r_vals, h_vals))
@@ -7887,6 +8029,9 @@ class RipleyDialog(QDialog):
             self.accept()
 
         except Exception as e:
+            import traceback
+            print(traceback.format_exc())
+
             QMessageBox.critical(
                 self,
                 "Error:",
@@ -8629,10 +8774,10 @@ class ResizeDialog(QDialog):
             else:
                 new_shape = tuple(int(dim * factor) for dim, factor in zip(array_shape, resize))
                 
-            if any(dim < 1 for dim in new_shape):
-                QMessageBox.critical(self, "Error", f"Resize would result in invalid dimensions: {new_shape}")
-                self.reset_fields()
-                return
+            #if any(dim < 1 for dim in new_shape):
+                #QMessageBox.critical(self, "Error", f"Resize would result in invalid dimensions: {new_shape}")
+                #self.reset_fields()
+                #return
 
             cubic = self.cubic.isChecked()
             order = 3 if cubic else 0
@@ -8706,7 +8851,12 @@ class ResizeDialog(QDialog):
                     centroids = copy.deepcopy(my_network.node_centroids)
                     if isinstance(resize, (int, float)):
                         for item in my_network.node_centroids:
-                            centroids[item] = np.round((my_network.node_centroids[item]) * resize)
+                            try:
+                                centroids[item] = np.round((my_network.node_centroids[item]) * resize)
+                            except:
+                                temp = np.array(my_network.node_centroids[item])
+                                centroids[item] = np.round((temp) * resize)
+
                     else:
                         for item in my_network.node_centroids:
                             centroids[item][0] = int(np.round((my_network.node_centroids[item][0]) * resize[0]))
@@ -9149,6 +9299,12 @@ class ThresholdDialog(QDialog):
 
     def start_ml(self, GPU = False):
 
+        try:
+            print("Please select image to load into nodes channel for segmentation or press X if you already have the one you want. Note that this load may permit a color image in the nodes channel for segmentation purposes only, which is otherwise not allowed.")
+            self.parent().load_channel(0, color = True)
+        except:
+            pass
+
 
         if self.parent().channel_data[2] is not None or self.parent().channel_data[3] is not None or self.parent().highlight_overlay is not None:
             if self.confirm_machine_dialog():
@@ -9188,13 +9344,14 @@ class ThresholdDialog(QDialog):
 
 class ExcelotronManager(QObject):
     # Signal to emit when data is received from Excelotron
-    data_received = pyqtSignal(dict, str)  # dictionary, property_name
+    data_received = pyqtSignal(dict, str, bool)  # dictionary, property_name
     
     def __init__(self, parent=None):
         super().__init__(parent)
         self.excelotron_window = None
         self.last_data = None
         self.last_property = None
+        self.last_add = None
     
     def launch(self):
         """Launch the Excelotron window"""
@@ -9243,12 +9400,13 @@ class ExcelotronManager(QObject):
         is_open = self.excelotron_window is not None
         return is_open
     
-    def _on_data_exported(self, data_dict, property_name):
+    def _on_data_exported(self, data_dict, property_name, add):
         """Internal slot to handle data from Excelotron"""
         self.last_data = data_dict
         self.last_property = property_name
+        self.last_add = add
         # Re-emit the signal for parent to handle
-        self.data_received.emit(data_dict, property_name)
+        self.data_received.emit(data_dict, property_name, add)
     
     def _on_window_destroyed(self):
         """Handle when the Excelotron window is destroyed/closed"""
@@ -9256,213 +9414,215 @@ class ExcelotronManager(QObject):
     
     def get_last_data(self):
         """Get the last exported data"""
-        return self.last_data, self.last_property
+        return self.last_data, self.last_property, self.last_add
 
 class MachineWindow(QMainWindow):
 
     def __init__(self, parent=None, GPU = False):
         super().__init__(parent)
 
-        if self.parent().active_channel == 0:
-            if self.parent().channel_data[0] is not None:
-                try:
-                    active_data = self.parent().channel_data[0]
-                    act_channel = 0
-                except:
+        try:
+
+            if self.parent().active_channel == 0:
+                if self.parent().channel_data[0] is not None:
+                    try:
+                        active_data = self.parent().channel_data[0]
+                        act_channel = 0
+                    except:
+                        active_data = self.parent().channel_data[1]
+                        act_channel = 1
+                else:
                     active_data = self.parent().channel_data[1]
                     act_channel = 1
-            else:
-                active_data = self.parent().channel_data[1]
-                act_channel = 1
-
-        try:
-            array1 = np.zeros_like(active_data).astype(np.uint8)
-        except:
-            print("No data in nodes channel")
-            return
 
-        self.setWindowTitle("Threshold")
-        
-        # Create central widget and layout
-        central_widget = QWidget()
-        self.setCentralWidget(central_widget)
-        layout = QVBoxLayout(central_widget)
+            try:
+                if len(active_data.shape) == 3:
+                    array1 = np.zeros_like(active_data).astype(np.uint8)
+                elif len(active_data.shape) == 4:
+                    array1 = np.zeros_like(active_data)[:,:,:,0].astype(np.uint8)
+            except:
+                print("No data in nodes channel")
+                return
 
+            self.setWindowTitle("Threshold")
+            
+            # Create central widget and layout
+            central_widget = QWidget()
+            self.setCentralWidget(central_widget)
+            layout = QVBoxLayout(central_widget)
 
-        # Create form layout for inputs
-        form_layout = QFormLayout()
 
-        layout.addLayout(form_layout)
+            # Create form layout for inputs
+            form_layout = QFormLayout()
 
-        if self.parent().pen_button.isChecked(): #Disable the pen mode if the user is in it because the segmenter pen forks it
-            self.parent().pen_button.click()
-        self.parent().threed = False
-        self.parent().can = False
-        self.parent().last_change = None
+            layout.addLayout(form_layout)
 
-        self.parent().pen_button.setEnabled(False)
+            if self.parent().pen_button.isChecked(): #Disable the pen mode if the user is in it because the segmenter pen forks it
+                self.parent().pen_button.click()
+            self.parent().threed = False
+            self.parent().can = False
+            self.parent().last_change = None
 
-        array3 = np.zeros_like(active_data).astype(np.uint8)
-        self.parent().highlight_overlay = array3 #Clear this out for the segmenter to use
+            self.parent().pen_button.setEnabled(False)
 
-        self.parent().load_channel(2, array1, True)
-        # Enable the channel button
-        # Not exactly sure why we need all this but the channel buttons weren't loading like they normally do when load_channel() is called:
-        if not self.parent().channel_buttons[2].isEnabled():
-            self.parent().channel_buttons[2].setEnabled(True)
-            self.parent().channel_buttons[2].click()
-        self.parent().delete_buttons[2].setEnabled(True)
+            array3 = np.zeros_like(array1).astype(np.uint8)
+            self.parent().highlight_overlay = array3 #Clear this out for the segmenter to use
 
-        self.parent().base_colors[act_channel] = self.parent().color_dictionary['WHITE']
-        self.parent().base_colors[2] = self.parent().color_dictionary['LIGHT_GREEN']
+            self.parent().load_channel(2, array1, True)
+            # Enable the channel button
+            # Not exactly sure why we need all this but the channel buttons weren't loading like they normally do when load_channel() is called:
+            if not self.parent().channel_buttons[2].isEnabled():
+                self.parent().channel_buttons[2].setEnabled(True)
+                self.parent().channel_buttons[2].click()
+            self.parent().delete_buttons[2].setEnabled(True)
 
-        self.parent().update_display()
-        
-        # Set a reasonable default size for the window
-        self.setMinimumWidth(600)  # Increased to accommodate grouped buttons
-        self.setMinimumHeight(500)
+            if len(active_data.shape) == 3:
+                self.parent().base_colors[act_channel] = self.parent().color_dictionary['WHITE']
+            self.parent().base_colors[2] = self.parent().color_dictionary['LIGHT_GREEN']
 
-        # Create main layout container
-        main_widget = QWidget()
-        main_layout = QVBoxLayout(main_widget)
+            self.parent().update_display()
+            
+            # Set a reasonable default size for the window
+            self.setMinimumWidth(600)  # Increased to accommodate grouped buttons
+            self.setMinimumHeight(500)
+
+            # Create main layout container
+            main_widget = QWidget()
+            main_layout = QVBoxLayout(main_widget)
+
+            # Group 1: Drawing tools (Brush + Foreground/Background)
+            drawing_group = QGroupBox("Drawing Tools")
+            drawing_layout = QHBoxLayout()
+
+            # Brush button
+            self.brush_button = QPushButton("🖌️")
+            self.brush_button.setCheckable(True)
+            self.brush_button.setFixedSize(40, 40)
+            self.brush_button.clicked.connect(self.toggle_brush_mode)
+            self.brush_button.click()
+
+            # Foreground/Background buttons in their own horizontal layout
+            fb_layout = QHBoxLayout()
+            self.fore_button = QPushButton("Foreground")
+            self.fore_button.setCheckable(True)
+            self.fore_button.setChecked(True)
+            self.fore_button.clicked.connect(self.toggle_foreground)
 
-        # Group 1: Drawing tools (Brush + Foreground/Background)
-        drawing_group = QGroupBox("Drawing Tools")
-        drawing_layout = QHBoxLayout()
+            self.back_button = QPushButton("Background")
+            self.back_button.setCheckable(True)
+            self.back_button.setChecked(False)
+            self.back_button.clicked.connect(self.toggle_background)
 
-        # Brush button
-        self.brush_button = QPushButton("🖌️")
-        self.brush_button.setCheckable(True)
-        self.brush_button.setFixedSize(40, 40)
-        self.brush_button.clicked.connect(self.toggle_brush_mode)
-        self.brush_button.click()
+            fb_layout.addWidget(self.fore_button)
+            fb_layout.addWidget(self.back_button)
 
-        # Foreground/Background buttons in their own horizontal layout
-        fb_layout = QHBoxLayout()
-        self.fore_button = QPushButton("Foreground")
-        self.fore_button.setCheckable(True)
-        self.fore_button.setChecked(True)
-        self.fore_button.clicked.connect(self.toggle_foreground)
+            drawing_layout.addWidget(self.brush_button)
+            drawing_layout.addLayout(fb_layout)
+            drawing_group.setLayout(drawing_layout)
 
-        self.back_button = QPushButton("Background")
-        self.back_button.setCheckable(True)
-        self.back_button.setChecked(False)
-        self.back_button.clicked.connect(self.toggle_background)
+            # Group 2: Processing Options (GPU)
+            processing_group = QGroupBox("Processing Options")
+            processing_layout = QHBoxLayout()
 
-        fb_layout.addWidget(self.fore_button)
-        fb_layout.addWidget(self.back_button)
+            self.use_gpu = GPU
+            self.two = QPushButton("Train By 2D Slice Patterns")
+            self.two.setCheckable(True)
+            self.two.setChecked(False)
+            self.two.clicked.connect(self.toggle_two)
+            self.use_two = False
+            self.three = QPushButton("Train by 3D Patterns")
+            self.three.setCheckable(True)
+            self.three.setChecked(True)
+            self.three.clicked.connect(self.toggle_three)
+            self.GPU = QPushButton("GPU")
+            self.GPU.setCheckable(True)
+            self.GPU.setChecked(False)
+            self.GPU.clicked.connect(self.toggle_GPU)
+            processing_layout.addWidget(self.GPU)
+            processing_layout.addWidget(self.two)
+            processing_layout.addWidget(self.three)
+            processing_group.setLayout(processing_layout)
+
+            # Group 3: Training Options
+            training_group = QGroupBox("Training")
+            training_layout = QHBoxLayout()
+            train_quick = QPushButton("Train Quick Model (When Good SNR)")
+            train_quick.clicked.connect(lambda: self.train_model(speed=True))
+            train_detailed = QPushButton("Train Detailed Model (For Morphology)")
+            train_detailed.clicked.connect(lambda: self.train_model(speed=False))
+            save = QPushButton("Save Model")
+            save.clicked.connect(self.save_model)
+            load = QPushButton("Load Model")
+            load.clicked.connect(self.load_model)
+            training_layout.addWidget(train_quick)
+            training_layout.addWidget(train_detailed)
+            training_layout.addWidget(save)
+            training_layout.addWidget(load)
+            training_group.setLayout(training_layout)
+
+            # Group 4: Segmentation Options
+            segmentation_group = QGroupBox("Segmentation")
+            segmentation_layout = QHBoxLayout()
+            seg_button = QPushButton("Preview Segment")
+            self.seg_button = seg_button
+            seg_button.clicked.connect(self.start_segmentation)
+            self.pause_button = QPushButton("▶/⏸️")
+            self.pause_button.setFixedSize(40, 40)
+            self.pause_button.clicked.connect(self.toggle_segment)
+            self.lock_button = QPushButton("🔒 Memory lock - (Prioritize RAM)")
+            self.lock_button.setCheckable(True)
+            self.lock_button.setChecked(True)
+            self.lock_button.clicked.connect(self.toggle_lock)
+            self.mem_lock = True
+            full_button = QPushButton("Segment All")
+            full_button.clicked.connect(self.segment)
+            segmentation_layout.addWidget(seg_button)
+            segmentation_layout.addWidget(self.pause_button)   # <--- for some reason the segmenter preview is still running even when killed, may be regenerating itself somewhere. May or may not actually try to resolve this because this feature isnt that necessary.
+            #segmentation_layout.addWidget(self.lock_button)   # Also turned this off
+            segmentation_layout.addWidget(full_button)
+            segmentation_group.setLayout(segmentation_layout)
+
+            # Add all groups to main layout
+            main_layout.addWidget(drawing_group)
+            if not GPU:
+                main_layout.addWidget(processing_group)
+            main_layout.addWidget(training_group)
+            main_layout.addWidget(segmentation_group)
+
+            # Set the main widget as the central widget
+            self.setCentralWidget(main_widget)
+
+            self.trained = False
+            self.previewing = False
 
-        drawing_layout.addWidget(self.brush_button)
-        drawing_layout.addLayout(fb_layout)
-        drawing_group.setLayout(drawing_layout)
+            if not GPU:
+                self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=False)
+            else:
+                self.segmenter = seg_GPU.InteractiveSegmenter(active_data)
 
-        # Group 2: Processing Options (GPU)
-        processing_group = QGroupBox("Processing Options")
-        processing_layout = QHBoxLayout()
-
-        self.use_gpu = GPU
-        self.two = QPushButton("Train By 2D Slice Patterns")
-        self.two.setCheckable(True)
-        self.two.setChecked(False)
-        self.two.clicked.connect(self.toggle_two)
-        self.use_two = False
-        self.three = QPushButton("Train by 3D Patterns")
-        self.three.setCheckable(True)
-        self.three.setChecked(True)
-        self.three.clicked.connect(self.toggle_three)
-        self.GPU = QPushButton("GPU")
-        self.GPU.setCheckable(True)
-        self.GPU.setChecked(False)
-        self.GPU.clicked.connect(self.toggle_GPU)
-        processing_layout.addWidget(self.GPU)
-        processing_layout.addWidget(self.two)
-        processing_layout.addWidget(self.three)
-        processing_group.setLayout(processing_layout)
+            self.segmentation_worker = None
 
-        # Group 3: Training Options
-        training_group = QGroupBox("Training")
-        training_layout = QHBoxLayout()
-        train_quick = QPushButton("Train Quick Model")
-        train_quick.clicked.connect(lambda: self.train_model(speed=True))
-        train_detailed = QPushButton("Train More Detailed Model")
-        train_detailed.clicked.connect(lambda: self.train_model(speed=False))
-        save = QPushButton("Save Model")
-        save.clicked.connect(self.save_model)
-        load = QPushButton("Load Model")
-        load.clicked.connect(self.load_model)
-        training_layout.addWidget(train_quick)
-        training_layout.addWidget(train_detailed)
-        training_layout.addWidget(save)
-        training_layout.addWidget(load)
-        training_group.setLayout(training_layout)
-
-        # Group 4: Segmentation Options
-        segmentation_group = QGroupBox("Segmentation")
-        segmentation_layout = QHBoxLayout()
-        seg_button = QPushButton("Preview Segment")
-        self.seg_button = seg_button
-        seg_button.clicked.connect(self.start_segmentation)
-        self.pause_button = QPushButton("▶/⏸️")
-        self.pause_button.clicked.connect(self.pause)
-        self.lock_button = QPushButton("🔒 Memory lock - (Prioritize RAM)")
-        self.lock_button.setCheckable(True)
-        self.lock_button.setChecked(True)
-        self.lock_button.clicked.connect(self.toggle_lock)
-        self.mem_lock = True
-        full_button = QPushButton("Segment All")
-        full_button.clicked.connect(self.segment)
-        segmentation_layout.addWidget(seg_button)
-        #segmentation_layout.addWidget(self.pause_button)   # <--- for some reason the segmenter preview is still running even when killed, may be regenerating itself somewhere. May or may not actually try to resolve this because this feature isnt that necessary.
-        #segmentation_layout.addWidget(self.lock_button)   # Also turned this off
-        segmentation_layout.addWidget(full_button)
-        segmentation_group.setLayout(segmentation_layout)
-
-        # Add all groups to main layout
-        main_layout.addWidget(drawing_group)
-        if not GPU:
-            main_layout.addWidget(processing_group)
-        main_layout.addWidget(training_group)
-        main_layout.addWidget(segmentation_group)
-
-        # Set the main widget as the central widget
-        self.setCentralWidget(main_widget)
+            self.fore_button.click()
+            self.fore_button.click()
 
-        self.trained = False
-        self.previewing = False
+        except:
+            return
 
-        if not GPU:
-            self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=False)
-        else:
-            self.segmenter = seg_GPU.InteractiveSegmenter(active_data)
+    def toggle_segment(self):
 
-        self.segmentation_worker = None
+        if self.segmentation_worker is not None:
+            if not self.segmentation_worker._paused:
+                self.segmentation_worker.pause()
+                print("Segmentation Worker Paused")
+            elif self.segmentation_worker._paused:
+                self.segmentation_worker.resume()
+                print("Segmentation Worker Resuming")
 
-        self.fore_button.click()
-        self.fore_button.click()
 
     def toggle_lock(self):
 
         self.mem_lock = self.lock_button.isChecked()
 
-    def pause(self):
-
-        if self.segmentation_worker is not None:
-            try:
-                print("Pausing segmenter")
-                self.previewing = False
-                self.segmentation_finished
-                del self.segmentation_worker
-                self.segmentation_worker = None
-            except:
-                pass
-
-        else:
-            try:
-                print("Restarting segmenter")
-                self.previewing = True
-                self.start_segmentation
-            except:
-                pass
 
     def save_model(self):
 
@@ -9484,6 +9644,9 @@ class MachineWindow(QMainWindow):
 
         except Exception as e:
             print(f"Error saving model: {e}")
+            import traceback
+            traceback.print_exc()
+
 
     def load_model(self):
 
@@ -9579,7 +9742,23 @@ class MachineWindow(QMainWindow):
     def toggle_brush_mode(self):
         """Toggle brush mode on/off"""
         self.parent().brush_mode = self.brush_button.isChecked()
+
         if self.parent().brush_mode:
+
+            self.parent().pm = painting.PaintManager(parent = self.parent())
+
+            # Start virtual paint session
+            # Get current zoom to preserve it
+            current_xlim = self.parent().ax.get_xlim() if hasattr(self, 'ax') and self.ax.get_xlim() != (0, 1) else None
+            current_ylim = self.parent().ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
+
+            if self.parent().pen_button.isChecked():
+                channel = self.parent().active_channel
+            else:
+                channel = 2
+
+            self.parent().pm.initiate_paint_session(channel, current_xlim, current_ylim)
+
             self.parent().pan_button.setChecked(False)
             self.parent().zoom_button.setChecked(False)
             if self.parent().pan_mode:
@@ -9606,12 +9785,13 @@ class MachineWindow(QMainWindow):
         self.kill_segmentation()
         # Wait a bit for cleanup
         time.sleep(0.1)
-        if not self.use_two:
-            self.previewing = False
+
+        self.previewing = True
         try:
             try:
                 self.segmenter.train_batch(self.parent().channel_data[2], speed = speed, use_gpu = self.use_gpu, use_two = self.use_two, mem_lock = self.mem_lock)
                 self.trained = True
+                self.start_segmentation()
             except Exception as e:
                 print("Error training. Perhaps you forgot both foreground and background markers? I need both!")
                 import traceback
@@ -9627,6 +9807,8 @@ class MachineWindow(QMainWindow):
 
     def start_segmentation(self):
 
+        self.parent().static_background = None
+
         self.kill_segmentation()
         time.sleep(0.1)
 
@@ -9635,12 +9817,10 @@ class MachineWindow(QMainWindow):
         else:
             print("Beginning new segmentation...")
 
-
-        if self.parent().active_channel == 0:
-            if self.parent().channel_data[0] is not None:
-                active_data = self.parent().channel_data[0]
-            else:
-                active_data = self.parent().channel_data[1]
+        if self.parent().channel_data[2] is not None:
+            active_data = self.parent().channel_data[2]
+        else:
+            active_data = self.parent().channel_data[0]
 
         array3 = np.zeros_like(active_data).astype(np.uint8)
         self.parent().highlight_overlay = array3 #Clear this out for the segmenter to use
@@ -9649,8 +9829,7 @@ class MachineWindow(QMainWindow):
             return
         else:
             self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu, self.use_two, self.previewing, self, self.mem_lock)
-            self.segmentation_worker.chunk_processed.connect(self.update_display)  # Just update display
-            self.segmentation_worker.finished.connect(self.segmentation_finished)
+            self.segmentation_worker.chunk_processed.connect(lambda: self.update_display(skip_paint_reinit = True))  # Just update display
             current_xlim = self.parent().ax.get_xlim()
             current_ylim = self.parent().ax.get_ylim()
             try:
@@ -9703,7 +9882,7 @@ class MachineWindow(QMainWindow):
             
         return changed
 
-    def update_display(self):
+    def update_display(self, skip_paint_reinit = False):
         if not hasattr(self, '_last_update'):
             self._last_update = 0
         
@@ -9713,8 +9892,6 @@ class MachineWindow(QMainWindow):
 
         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
@@ -9731,71 +9908,40 @@ class MachineWindow(QMainWindow):
 
                 if not self.parent().painting:
                     # Only update if view limits are valid
-                    self.parent().update_display(preserve_zoom=(current_xlim, current_ylim))
+                    self.parent().update_display(preserve_zoom=(current_xlim, current_ylim), skip_paint_reinit = skip_paint_reinit)
+
+                    if self.parent().brush_mode:
+                        current_xlim = self.parent().ax.get_xlim() if hasattr(self, 'ax') and self.ax.get_xlim() != (0, 1) else None
+                        current_ylim = self.parent().ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
+
+                        if self.parent().pen_button.isChecked():
+                            channel = self.parent().active_channel
+                        else:
+                            channel = 2
 
+                        self.parent().pm.initiate_paint_session(channel, current_xlim, current_ylim)
                     
                     self._last_update = current_time
             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()
+        try:
+            # Clear any processing flags in the segmenter
+            if hasattr(self.segmenter, '_currently_processing'):
+                self.segmenter._currently_processing = None
                 
-            except Exception as e:
-                print(f"Error in poke_segmenter: {e}")
-                import traceback
-                traceback.print_exc()
-
-    def segmentation_finished(self):
-        
-        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 = {}
+            # Force regenerating the worker
+            if self.segmentation_worker is not None:
+                self.kill_segmentation()
                 
-            current_time = time.time()
+            time.sleep(0.2)
+            self.start_segmentation()
             
-            # 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
-                
-                if 0 in self.parent().highlight_overlay[current_z, :, :]:
-                    # Create a new worker after a brief delay
-                    QTimer.singleShot(500, self.start_segmentation)
-
+        except Exception as e:
+            print(f"Error in poke_segmenter: {e}")
+            import traceback
+            traceback.print_exc()
 
 
     def kill_segmentation(self):
@@ -9829,11 +9975,10 @@ class MachineWindow(QMainWindow):
 
             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]
-                else:
-                    active_data = self.parent().channel_data[1]
+            if self.parent().channel_data[2] is not None:
+                active_data = self.parent().channel_data[2]
+            else:
+                active_data = self.parent().channel_data[0]
 
             array3 = np.zeros_like(active_data).astype(np.uint8)
             self.parent().highlight_overlay = array3 #Clear this out for the segmenter to use
@@ -9844,6 +9989,8 @@ class MachineWindow(QMainWindow):
                 self.parent().highlight_overlay = self.segmenter.segment_volume(array = self.parent().highlight_overlay)
             except Exception as e:
                 print(f"Error segmenting (Perhaps retrain the model...): {e}")
+                import traceback
+                traceback.print_exc()
                 return
 
             # Clean up when done
@@ -9880,6 +10027,12 @@ class MachineWindow(QMainWindow):
                     # Kill the segmentation thread and wait for it to finish
                     self.kill_segmentation()
                     time.sleep(0.2)  # Give additional time for cleanup
+
+                    try:
+                        self.parent().channel_data[0] = self.parent().reduce_rgb_dimension(self.parent().channel_data[0], 'weight')
+                        self.update_display()
+                    except:
+                        pass
                     
                     self.parent().machine_window = None
                 else:
@@ -9940,59 +10093,26 @@ class SegmentationWorker(QThread):
             
             # 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):
-                    # Check for pause/stop before processing each chunk
-                    self._check_pause()
-                    if self._stop:
-                        break
-                        
-                    if foreground == None and background == None:
-                        self.get_poked()
 
-                    if self._stop:
-                        break
-                    
-                    # 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
+            # Original 3D approach
+            for foreground_coords, background_coords in self.segmenter.segment_volume_realtime(gpu=self.use_gpu):
+                # Check for pause/stop before processing each chunk
+                self._check_pause()
+                if self._stop:
+                    break
                 
-            else:
-                # Original 3D approach
-                for foreground_coords, background_coords in self.segmenter.segment_volume_realtime(gpu=self.use_gpu):
-                    # Check for pause/stop before processing each chunk
-                    self._check_pause()
-                    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 
+                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()
             
@@ -10780,7 +10900,7 @@ class MaskDialog(QDialog):
 
 class CropDialog(QDialog):
 
-    def __init__(self, parent=None):
+    def __init__(self, parent=None, args = None):
 
         try:
 
@@ -10788,18 +10908,26 @@ class CropDialog(QDialog):
             self.setWindowTitle("Crop Image (Will transpose any centroids)?")
             self.setModal(True)
 
+            if args is None:
+                xmin = 0
+                xmax = self.parent().shape[2]
+                ymin = 0
+                ymax = self.parent().shape[1]
+            else:
+                xmin, xmax, ymin, ymax = args
+
             layout = QFormLayout(self)
 
-            self.xmin = QLineEdit("0")
+            self.xmin = QLineEdit(f"{xmin}")
             layout.addRow("X Min", self.xmin)
 
-            self.xmax = QLineEdit(f"{self.parent().shape[2]}")
+            self.xmax = QLineEdit(f"{xmax}")
             layout.addRow("X Max", self.xmax)
 
-            self.ymin = QLineEdit("0")
+            self.ymin = QLineEdit(f"{ymin}")
             layout.addRow("Y Min", self.ymin)
 
-            self.ymax = QLineEdit(f"{self.parent().shape[1]}")
+            self.ymax = QLineEdit(f"{ymax}")
             layout.addRow("Y Max", self.ymax)
 
             self.zmin = QLineEdit("0")
@@ -10853,10 +10981,16 @@ class CropDialog(QDialog):
                     transformed = centroids - np.array([zmin, ymin, xmin])
                     transformed = transformed.astype(int)
                     
-                    # Boolean mask for valid coordinates
-                    valid_mask = ((transformed >= 0) & 
-                                  (transformed <= np.array([zmax, ymax, xmax]))).all(axis=1)
+                    # Create upper bounds array with same shape
+                    upper_bounds = np.array([zmax - zmin, ymax - ymin, xmax - xmin])
                     
+                    # Boolean mask for valid coordinates - check each dimension separately
+                    z_valid = (transformed[:, 0] >= 0) & (transformed[:, 0] <= upper_bounds[0])
+                    y_valid = (transformed[:, 1] >= 0) & (transformed[:, 1] <= upper_bounds[1])
+                    x_valid = (transformed[:, 2] >= 0) & (transformed[:, 2] <= upper_bounds[2])
+                    
+                    valid_mask = z_valid & y_valid & x_valid
+
                     # Rebuild dictionary with only valid entries
                     my_network.node_centroids = {
                         nodes[int(i)]: [int(transformed[i, 0]), int(transformed[i, 1]), int(transformed[i, 2])]
@@ -10865,6 +10999,15 @@ class CropDialog(QDialog):
                     
                     self.parent().format_for_upperright_table(my_network.node_centroids, 'NodeID', ['Z', 'Y', 'X'], 'Node Centroids')
 
+                if my_network.node_identities is not None:
+                    new_idens = {}
+                    for node, iden in my_network.node_identities.items():
+                        if node in my_network.node_centroids:
+                            new_idens[node] = iden
+                    my_network.node_identities = new_idens
+
+                    self.parent().format_for_upperright_table(my_network.node_identities, 'NodeID', 'Identity', 'Node Identities')
+
             except Exception as e:
 
                 print(f"Error transposing node centroids: {e}")
@@ -11058,6 +11201,74 @@ class SkeletonizeDialog(QDialog):
                 f"Error running skeletonize: {str(e)}"
             )   
 
+class DistanceDialog(QDialog):
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.setWindowTitle("Compute Distance Transform (Applies xy and z scaling, set them to 1 if you want voxel correspondence)?")
+        self.setModal(True)
+        
+        layout = QFormLayout(self)
+
+        # Add Run button
+        run_button = QPushButton("Run")
+        run_button.clicked.connect(self.run)
+        layout.addRow(run_button)
+
+    def run(self):
+
+        try:
+
+            data = self.parent().channel_data[self.parent().active_channel]
+
+            data = sdl.compute_distance_transform_distance(data, sampling = [my_network.z_scale, my_network.xy_scale, my_network.xy_scale])
+
+            self.parent().load_channel(self.parent().active_channel, data, data = True)
+
+        except Exception as e:
+
+            print(f"Error: {e}")
+
+class GrayWaterDialog(QDialog):
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.setWindowTitle(f"Gray Watershed - Please segment out your background first (ie with intensity thresholding) or this will not work correctly. \nAt the moment, this is designed for similarly sized objects. Having mixed large/small objects may not work correctly.")
+        self.setModal(True)
+        
+        layout = QFormLayout(self)
+
+        self.min_peak_distance = QLineEdit("1")
+        layout.addRow("Minimum Peak Distance (To any other peak - Recommended) (This is true voxel distance here)", self.min_peak_distance)
+
+        # Minimum Intensity
+        self.min_intensity = QLineEdit("")
+        layout.addRow("Minimum Peak Intensity (Optional):", self.min_intensity)
+
+        # Add Run button
+        run_button = QPushButton("Run Watershed")
+        run_button.clicked.connect(self.run_watershed)
+        layout.addRow(run_button)
+
+    def run_watershed(self):
+
+        try:
+
+            min_intensity = float(self.min_intensity.text()) if self.min_intensity.text().strip() else None
+
+            min_peak_distance = int(self.min_peak_distance.text()) if self.min_peak_distance.text().strip() else 1
+
+            data = self.parent().channel_data[self.parent().active_channel]
+
+            data = n3d.gray_watershed(data, min_peak_distance, min_intensity)
+
+            self.parent().load_channel(self.parent().active_channel, data, data = True)
+
+            self.accept()
+
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+
 
 class WatershedDialog(QDialog):
     def __init__(self, parent=None):
@@ -11084,10 +11295,14 @@ class WatershedDialog(QDialog):
         except:
             self.default = 0.05
 
+        # Smallest radius (empty by default)
+        self.smallest_rad = QLineEdit()
+        self.smallest_rad.setPlaceholderText("Leave empty for None")
+        layout.addRow(f"Smallest Radius (Objects any smaller may get thresholded out - this value always overrides below 'proportion' param). \n Somewhat more intuitive param then below, use a conservative value a bit smaller than your smallest object's radius:", self.smallest_rad)
         
         # Proportion (default 0.1)
         self.proportion = QLineEdit(f"{self.default}")
-        layout.addRow("Proportion:", self.proportion)
+        layout.addRow(f"Proportion (0-1) of distance transform value set [ie unique elements] to exclude (ie 0.2 = 20% of the set of all values of the distance transform get excluded).\n Essentially, vals closer to 0 are less likely to split objects but also won't kick out small objects from the output, vals slightly further from 0 will split more aggressively, but vals closer to 1 become unstable, leading to objects being evicted or labelling errors. \nRecommend something between 0.05 and 0.4, but it depends on the data (Or just enter a smallest radius above to avoid using this). \nWill tell you in command window what equivalent 'smallest radius' this is):", self.proportion)
         
         # GPU checkbox (default True)
         self.gpu = QPushButton("GPU")
@@ -11095,10 +11310,6 @@ class WatershedDialog(QDialog):
         self.gpu.setChecked(False)
         layout.addRow("Use GPU:", self.gpu)
         
-        # Smallest radius (empty by default)
-        self.smallest_rad = QLineEdit()
-        self.smallest_rad.setPlaceholderText("Leave empty for None")
-        layout.addRow("Smallest Radius:", self.smallest_rad)
         
         # Predownsample (empty by default)
         self.predownsample = QLineEdit()
@@ -11106,11 +11317,11 @@ class WatershedDialog(QDialog):
         layout.addRow("Kernel Obtainment GPU Downsample:", self.predownsample)
         
         # Predownsample2 (empty by default)
-        self.predownsample2 = QLineEdit()
-        self.predownsample2.setPlaceholderText("Leave empty for None")
-        layout.addRow("Smart Label GPU Downsample:", self.predownsample2)
+        #self.predownsample2 = QLineEdit()
+        #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'."))
+        #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
@@ -11147,7 +11358,7 @@ class WatershedDialog(QDialog):
             # Get predownsample2 (None if empty)
             try:
                 predownsample2 = float(self.predownsample2.text()) if self.predownsample2.text() else None
-            except ValueError:
+            except:
                 predownsample2 = None
             
             # Get the active channel data from parent
@@ -11325,7 +11536,22 @@ class CentroidNodeDialog(QDialog):
 
             if mode == 0:
 
-                my_network.nodes = my_network.centroid_array()
+                try:
+                    shape = my_network.nodes.shape
+
+                except:
+                    try:
+                        shape = my_network.edges.shape
+                    except:
+                        try:
+                            shape = my_network.network_overlay.shape
+                        except:
+                            try:
+                                shape = my_network.id_overlay.shape
+                            except:
+                                shape = None
+
+                my_network.nodes = my_network.centroid_array(shape = shape)
 
             else:
 
@@ -11439,7 +11665,7 @@ class GenNodesDialog(QDialog):
         
         # Component dilation
         self.comp_dil = QLineEdit("0")
-        opt_layout.addWidget(QLabel("Voxel distance to merge nearby nodes (Compensates for multi-branch regions):"), 1, 0)
+        opt_layout.addWidget(QLabel("Amount to expand nodes (Merges nearby nodes, say if they are overassigned, good for broader branch breaking):"), 1, 0)
         opt_layout.addWidget(self.comp_dil, 1, 1)
         
         opt_group.setLayout(opt_layout)
@@ -11573,11 +11799,11 @@ class GenNodesDialog(QDialog):
 
 class BranchDialog(QDialog):
 
-    def __init__(self, parent=None):
+    def __init__(self, parent=None, called = False):
         super().__init__(parent)
         self.setWindowTitle("Label Branches (of edges)")
         self.setModal(True)
-        
+
         # Main layout
         main_layout = QVBoxLayout(self)
         
@@ -11610,7 +11836,10 @@ class BranchDialog(QDialog):
 
         self.fix3 = QPushButton("Split Nontouching Branches?")
         self.fix3.setCheckable(True)
-        self.fix3.setChecked(True)
+        if called:
+            self.fix3.setChecked(True)
+        else:
+            self.fix3.setChecked(False)
         correction_layout.addWidget(QLabel("Split Nontouching Branches? (Useful if branch pruning - may want to threshold out small, split branches after): "), 4, 0)
         correction_layout.addWidget(self.fix3, 4, 1)
         
@@ -12137,6 +12366,8 @@ class CentroidDialog(QDialog):
 
         try:
 
+            print("Calculating centroids...")
+
             chan = self.mode_selector.currentIndex()
 
             # Get directory (None if empty)
@@ -12499,6 +12730,9 @@ class CalcAllDialog(QDialog):
 
             
         except Exception as e:
+            import traceback
+            print(traceback.format_exc())
+
             QMessageBox.critical(
                 self,
                 "Error",
@@ -12506,6 +12740,7 @@ class CalcAllDialog(QDialog):
             )
 
 
+
 class ProxDialog(QDialog):
     def __init__(self, parent=None):
         super().__init__(parent)
@@ -12744,11 +12979,6 @@ class ProxDialog(QDialog):
 
 
 
-
-
-
-
-
 # Initiating this program from the script line:
 
 def run_gui():