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

nettracer3d/nettracer_gui.py

@@ -35,6 +35,7 @@ import queue
 from threading import Lock
 from scipy import ndimage
 import os
+from . import painting
 
 
 
@@ -210,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)
@@ -454,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
         
@@ -2170,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
@@ -2184,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)
@@ -2195,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)
@@ -2210,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()
@@ -2224,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()
@@ -2443,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()
@@ -2486,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."""
@@ -2533,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)
@@ -2597,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):
 
@@ -2823,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
+            # Update display with virtual paint (super fast)
+            self.pm.update_virtual_paint_display()
 
-    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']
-                    )
-
-    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"""
@@ -3322,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()
 
             
 
@@ -3675,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")
@@ -3779,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:
@@ -4555,6 +4367,7 @@ class ImageViewerWindow(QMainWindow):
                     )
 
             self.last_load = directory
+            
 
             if directory != "":
 
@@ -4817,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:
         --------
@@ -4831,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':
@@ -4840,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)
 
@@ -4863,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:
@@ -4940,12 +4757,12 @@ class ImageViewerWindow(QMainWindow):
             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
 
@@ -5295,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:
@@ -5321,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:
+                if not continue_paint:
+                    self.static_background = 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
+                    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]
@@ -5572,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'):
@@ -5744,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"""
         
@@ -7264,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):
@@ -7624,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()
@@ -9503,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():
@@ -9619,206 +9421,208 @@ 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):
 
@@ -9840,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):
 
@@ -9935,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:
@@ -9962,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
@@ -9983,6 +9807,8 @@ class MachineWindow(QMainWindow):
 
     def start_segmentation(self):
 
+        self.parent().static_background = None
+
         self.kill_segmentation()
         time.sleep(0.1)
 
@@ -9991,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
@@ -10005,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:
@@ -10059,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
         
@@ -10069,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
@@ -10087,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):
@@ -10185,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
@@ -10200,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
@@ -10236,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:
@@ -10296,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()