nettracer3d 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

nettracer3d/nettracer_gui.py

@@ -16,13 +16,15 @@ from nettracer3d import nettracer as n3d
 from nettracer3d import smart_dilate as sdl
 from nettracer3d import proximity as pxt
 from matplotlib.colors import LinearSegmentedColormap
+from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
 import pandas as pd
-from PyQt6.QtGui import (QFont, QCursor, QColor)
+from PyQt6.QtGui import (QFont, QCursor, QColor, QPixmap, QPainter, QPen)
 import tifffile
 import copy
 import multiprocessing as mp
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
+from nettracer3d import segmenter
 
 
 class ImageViewerWindow(QMainWindow):
@@ -105,11 +107,21 @@ class ImageViewerWindow(QMainWindow):
         self.zoom_mode = False
         self.original_xlim = None
         self.original_ylim = None
+        self.zoom_changed = False
 
         # Pan mode state
         self.pan_mode = False
         self.panning = False
         self.pan_start = None
+
+        #For ML segmenting mode
+        self.brush_mode = False
+        self.painting = False
+        self.foreground = True
+        self.machine_window = None
+        self.brush_size = 1  # Start with 1 pixel
+        self.min_brush_size = 1
+        self.max_brush_size = 10
         
         # Store brightness/contrast values for each channel
         self.channel_brightness = [{
@@ -176,6 +188,9 @@ class ImageViewerWindow(QMainWindow):
         buttons_layout.addWidget(self.pan_button)
 
         control_layout.addWidget(buttons_widget)
+
+        self.preview = False #Whether in preview mode or not
+        self.targs = None #Targets for preview mode
                 
         # Create channel buttons
         self.channel_buttons = []
@@ -236,6 +251,8 @@ class ImageViewerWindow(QMainWindow):
         self.ax = self.figure.add_subplot(111)
         left_layout.addWidget(self.canvas)
 
+        self.canvas.mpl_connect('scroll_event', self.on_mpl_scroll)
+
         
         left_layout.addWidget(control_panel)
 
@@ -344,6 +361,7 @@ class ImageViewerWindow(QMainWindow):
         self.canvas.mpl_connect('button_press_event', self.on_mouse_press)
         self.canvas.mpl_connect('button_release_event', self.on_mouse_release)
         self.canvas.mpl_connect('motion_notify_event', self.on_mouse_move)
+
         #self.canvas.mpl_connect('button_press_event', self.on_mouse_click)
 
         # Initialize measurement points tracking
@@ -396,7 +414,7 @@ class ImageViewerWindow(QMainWindow):
             self.slice_slider.setValue(new_value)
         
 
-    def create_highlight_overlay(self, node_indices=None, edge_indices=None, overlay1_indices = None, overlay2_indices = None):
+    def create_highlight_overlay(self, node_indices=None, edge_indices=None, overlay1_indices = None, overlay2_indices = None, bounds = False):
         """
         Create a binary overlay highlighting specific nodes and/or edges using parallel processing.
         
@@ -410,6 +428,12 @@ class ImageViewerWindow(QMainWindow):
             mask = np.isin(chunk_data, indices_to_check)
             return mask * 255
 
+        def process_chunk_bounds(chunk_data, indices_to_check):
+            """Process a single chunk of the array to create highlight mask"""
+
+            mask = (chunk_data >= indices_to_check[0]) & (chunk_data <= indices_to_check[1])
+            return mask * 255
+
         if node_indices is not None:
             if 0 in node_indices:
                 node_indices.remove(0)
@@ -453,7 +477,7 @@ class ImageViewerWindow(QMainWindow):
         num_cores = mp.cpu_count()
         
         # Calculate chunk size along y-axis
-        chunk_size = full_shape[0] // num_cores
+        chunk_size = full_shape[1] // num_cores
         if chunk_size < 1:
             chunk_size = 1
         
@@ -463,25 +487,32 @@ class ImageViewerWindow(QMainWindow):
                 
             # Create chunks
             chunks = []
-            for i in range(0, array_shape[0], chunk_size):
-                end = min(i + chunk_size, array_shape[0])
-                chunks.append(channel_data[i:end])
+            for i in range(0, array_shape[1], chunk_size):
+                end = min(i + chunk_size, array_shape[1])
+                chunks.append(channel_data[:, i:end, :])
                 
             # Process chunks in parallel using ThreadPoolExecutor
-            process_func = partial(process_chunk, indices_to_check=indices)
+            if not bounds:
+                process_func = partial(process_chunk, indices_to_check=indices)
+            else:
+                if len(indices) == 1:
+                    indices.insert(0, 0)
+                process_func = partial(process_chunk_bounds, indices_to_check=indices)
+
             
             with ThreadPoolExecutor(max_workers=num_cores) as executor:
                 chunk_results = list(executor.map(process_func, chunks))
                 
             # Reassemble the chunks
-            return np.vstack(chunk_results)
+            return np.concatenate(chunk_results, axis=1)
         
         # Process nodes and edges in parallel using multiprocessing
-        with ThreadPoolExecutor(max_workers=2) as executor:
+        with ThreadPoolExecutor(max_workers=num_cores) as executor:
             future_nodes = executor.submit(process_channel, self.channel_data[0], node_indices, full_shape)
             future_edges = executor.submit(process_channel, self.channel_data[1], edge_indices, full_shape)
             future_overlay1 = executor.submit(process_channel, self.channel_data[2], overlay1_indices, full_shape)
             future_overlay2 = executor.submit(process_channel, self.channel_data[3], overlay2_indices, full_shape)
+
             
             # Get results
             node_overlay = future_nodes.result()
@@ -502,6 +533,84 @@ class ImageViewerWindow(QMainWindow):
         # Update display
         self.update_display(preserve_zoom=(current_xlim, current_ylim))
 
+    def create_highlight_overlay_slice(self, indices, bounds = False):
+
+
+        def process_chunk_bounds(chunk_data, indices_to_check):
+            """Process a single chunk of the array to create highlight mask"""
+            mask = (chunk_data >= indices_to_check[0]) & (chunk_data <= indices_to_check[1])
+            return mask * 255
+
+        def process_chunk(chunk_data, indices_to_check):
+            """Process a single chunk of the array to create highlight mask"""
+
+            mask = np.isin(chunk_data, indices_to_check)
+            return mask * 255
+
+        array = self.channel_data[self.active_channel]
+        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_slice = array[self.current_slice, :, :]
+        full_shape = array.shape
+        slice_shape = current_slice.shape
+
+        if self.highlight_overlay is None:
+
+            self.highlight_overlay = np.zeros(full_shape, dtype=np.uint8)
+
+        # Get number of CPU cores
+        num_cores = mp.cpu_count()
+        
+        # Calculate chunk size along y-axis
+        chunk_size = slice_shape[0] // num_cores
+        if chunk_size < 1:
+            chunk_size = 1
+        
+        def process_channel(channel_data, indices, array_shape):
+            if channel_data is None or not indices:
+                return None
+                
+            # Create chunks
+            chunks = []
+            for i in range(0, array_shape[0], chunk_size):
+                end = min(i + chunk_size, array_shape[0])
+                chunks.append(channel_data[i:end])
+                
+            # Process chunks in parallel using ThreadPoolExecutor
+            if not bounds:
+                process_func = partial(process_chunk, indices_to_check=indices)
+            else:
+                if len(indices) == 1:
+                    indices.insert(0, 0)
+                process_func = partial(process_chunk_bounds, indices_to_check=indices)
+
+            
+            with ThreadPoolExecutor(max_workers=num_cores) as executor:
+                chunk_results = list(executor.map(process_func, chunks))
+                
+            # Reassemble the chunks
+            return np.vstack(chunk_results)
+        
+        # Process nodes and edges in parallel using multiprocessing
+        with ThreadPoolExecutor(max_workers=num_cores) as executor:
+            future_highlight = executor.submit(process_channel, current_slice, indices, slice_shape)
+            
+            # Get results
+            overlay = future_highlight.result()
+
+        try:
+
+            self.highlight_overlay[self.current_slice, :, :] = overlay
+        except:
+            pass
+
+        # Update display
+        self.update_display(preserve_zoom=(current_xlim, current_ylim), called = True)
+
+
+
+        
 
 
 
@@ -1328,9 +1437,16 @@ class ImageViewerWindow(QMainWindow):
         if self.zoom_mode:
             self.pan_button.setChecked(False)
             self.pan_mode = False
+            self.brush_mode = False
+            if self.machine_window is not None:
+                self.machine_window.silence_button()
             self.canvas.setCursor(Qt.CursorShape.CrossCursor)
         else:
-            self.canvas.setCursor(Qt.CursorShape.ArrowCursor)
+            if self.machine_window is None:
+                self.canvas.setCursor(Qt.CursorShape.ArrowCursor)
+            else:
+                self.machine_window.toggle_brush_button()
+
 
     def toggle_pan_mode(self):
         """Toggle pan mode on/off."""
@@ -1338,14 +1454,132 @@ class ImageViewerWindow(QMainWindow):
         if self.pan_mode:
             self.zoom_button.setChecked(False)
             self.zoom_mode = False
+            self.brush_mode = False
+            if self.machine_window is not None:
+                self.machine_window.silence_button()
             self.canvas.setCursor(Qt.CursorShape.OpenHandCursor)
         else:
-            self.canvas.setCursor(Qt.CursorShape.ArrowCursor)
+            if self.machine_window is None:
+                self.canvas.setCursor(Qt.CursorShape.ArrowCursor)
+            else:
+                self.machine_window.toggle_brush_button()
+
+
+
+    def on_mpl_scroll(self, event):
+        """Handle matplotlib canvas scroll events"""
+        #Wheel events
+        if self.brush_mode and event.inaxes == self.ax:
+            # Check if Ctrl is pressed
+            if event.guiEvent.modifiers() & Qt.ShiftModifier:
+                pass
+
+            elif event.guiEvent.modifiers() & Qt.ControlModifier:
+                # Brush size adjustment code...
+                step = 1 if event.button == 'up' else -1
+                new_size = self.brush_size + step
+                    
+                if new_size < self.min_brush_size:
+                    new_size = self.min_brush_size
+                elif new_size > self.max_brush_size:
+                    new_size = self.max_brush_size
+                    
+                self.brush_size = new_size
+                self.update_brush_cursor()
+                event.guiEvent.accept()
+                return
+
+        # General scrolling code outside the brush mode condition
+        step = 1 if event.button == 'up' else -1
+
+        if event.guiEvent.modifiers() & Qt.ShiftModifier:
+            if event.guiEvent.modifiers() & Qt.ControlModifier:
+                step = step * 3
+            if (self.current_slice + step) < 0 or (self.current_slice + step) > self.slice_slider.maximum():
+                return
+                
+            self.current_slice = self.current_slice + step
+            self.slice_slider.setValue(self.current_slice + step)
+
+            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
+            
+            self.update_display(preserve_zoom=(current_xlim, current_ylim))
+
+    def keyPressEvent(self, event):
+        if event.key() == Qt.Key_Z:
+            self.zoom_button.click()
+        if self.machine_window is not None:
+            if event.key() == Qt.Key_A:
+                self.machine_window.switch_foreground()
+
+
+    def update_brush_cursor(self):
+        """Update the cursor to show brush size"""
+        if not self.brush_mode:
+            return
+            
+        # Create a pixmap for the cursor
+        size = self.brush_size * 2 + 2  # Add padding for border
+        pixmap = QPixmap(size, size)
+        pixmap.fill(Qt.transparent)
+        
+        # Create painter for the pixmap
+        painter = QPainter(pixmap)
+        painter.setRenderHint(QPainter.RenderHint.Antialiasing) 
+        
+        # Draw circle
+        pen = QPen(Qt.white)
+        pen.setWidth(1)
+        painter.setPen(pen)
+        painter.setBrush(Qt.transparent)
+        painter.drawEllipse(1, 1, size-2, size-2)
+        
+        # Create cursor from pixmap
+        cursor = QCursor(pixmap)
+        self.canvas.setCursor(cursor)
+        
+        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 on_mouse_press(self, event):
         """Handle mouse press events."""
         if event.inaxes != self.ax:
             return
+
+        if event.button == 2:
+            self.pan_button.click()
+            return
                 
         if self.zoom_mode:
             # Handle zoom mode press
@@ -1364,6 +1598,11 @@ class ImageViewerWindow(QMainWindow):
                 
                 self.ax.set_xlim([xdata - x_range, xdata + x_range])
                 self.ax.set_ylim([ydata - y_range, ydata + y_range])
+
+                self.zoom_changed = True  # Flag that zoom has changed
+                
+                if not hasattr(self, 'zoom_changed'):
+                    self.zoom_changed = False
                 
             elif event.button == 3:  # Right click - zoom out
                 x_range = (current_xlim[1] - current_xlim[0])
@@ -1381,6 +1620,11 @@ class ImageViewerWindow(QMainWindow):
                 else:
                     self.ax.set_xlim(new_xlim)
                     self.ax.set_ylim(new_ylim)
+
+                    self.zoom_changed = False  # Flag that zoom has changed
+                    
+                    if not hasattr(self, 'zoom_changed'):
+                        self.zoom_changed = False
             
             self.canvas.draw()
                 
@@ -1388,6 +1632,50 @@ class ImageViewerWindow(QMainWindow):
             self.panning = True
             self.pan_start = (event.xdata, event.ydata)
             self.canvas.setCursor(Qt.CursorShape.ClosedHandCursor)
+
+        elif self.brush_mode:
+            if event.inaxes != self.ax:
+                return
+            
+
+            if event.button == 1 or event.button == 3:
+
+                if event.button == 3:
+                    self.erase = True
+                else:
+                    self.erase = False
+
+                self.painting = True
+                x, y = int(event.xdata), int(event.ydata)
+                self.last_paint_pos = (x, y)
+
+                if self.foreground:
+                    channel = 2
+                else:
+                    channel = 3
+                
+
+                # Paint at initial position
+                self.paint_at_position(x, y, self.erase, channel)
+                
+                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
+
+            
+                self.canvas.draw()
+                #self.update_display(preserve_zoom=(current_xlim, current_ylim))
+                self.restore_channels = []
+
+
+                for i in range(4):
+                    if i == channel:
+                        self.channel_visible[i] = True
+                    elif self.channel_data[i] is not None and self.channel_visible[i] == True:
+                        self.channel_visible[i] = False
+                        self.restore_channels.append(i)
+                self.update_display(preserve_zoom = (current_xlim, current_ylim), begin_paint = True)
+                self.update_display_slice(channel, preserve_zoom=(current_xlim, current_ylim))
+
         
         elif event.button == 3:  # Right click (for context menu)
             self.create_context_menu(event)
@@ -1397,12 +1685,32 @@ 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
+        else:
+            val = 255
+
+        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:
+                    self.channel_data[channel][self.current_slice][y, x] = val
+
     def on_mouse_move(self, event):
         """Handle mouse movement events."""
         if event.inaxes != self.ax:
             return
                 
-        if self.selection_start and not self.selecting and not self.pan_mode and not self.zoom_mode:
+        if self.selection_start and not self.selecting and not self.pan_mode and not self.zoom_mode and not self.brush_mode:
             # If mouse has moved more than a tiny amount while button is held, start selection
             if (abs(event.xdata - self.selection_start[0]) > 1 or 
                 abs(event.ydata - self.selection_start[1]) > 1):
@@ -1424,6 +1732,7 @@ class ImageViewerWindow(QMainWindow):
             self.canvas.draw()
 
         elif self.panning and self.pan_start is not None:
+
             # Calculate the movement
             dx = event.xdata - self.pan_start[0]
             dy = event.ydata - self.pan_start[1]
@@ -1459,6 +1768,39 @@ class ImageViewerWindow(QMainWindow):
             # Update pan start position
             self.pan_start = (event.xdata, event.ydata)
 
+        elif self.painting and self.brush_mode:
+            if event.inaxes != self.ax:
+                return
+                
+            x, y = int(event.xdata), int(event.ydata)
+
+            if self.foreground:
+                channel = 2
+            else:
+                channel = 3
+
+            
+            if self.channel_data[2] is not 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
+                height, width = self.channel_data[2][self.current_slice].shape
+                
+                if hasattr(self, 'last_paint_pos'):
+                    last_x, last_y = self.last_paint_pos
+                    points = self.get_line_points(last_x, last_y, x, y)
+                    
+                    # Paint at each point along the line
+                    for px, py in points:
+                        if 0 <= px < width and 0 <= py < height:
+                            self.paint_at_position(px, py, self.erase, channel)
+                
+                self.last_paint_pos = (x, y)
+                
+                self.canvas.draw()
+                #self.update_display(preserve_zoom=(current_xlim, current_ylim))
+                self.update_display_slice(channel, preserve_zoom=(current_xlim, current_ylim))
+
+
     def on_mouse_release(self, event):
         """Handle mouse release events."""
         if self.pan_mode:
@@ -1520,15 +1862,27 @@ class ImageViewerWindow(QMainWindow):
             elif not self.selecting and self.selection_start:  # If we had a click but never started selection
                 # Handle as a normal click
                 self.on_mouse_click(event)
+
             
             # Clean up
             self.selection_start = None
             self.selecting = False
+
+
             if self.selection_rect is not None:
                 self.selection_rect.remove()
                 self.selection_rect = None
                 self.canvas.draw()
 
+        if self.brush_mode:
+            self.painting = False
+            for i in self.restore_channels:
+                self.channel_visible[i] = True
+            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
+            self.update_display(preserve_zoom=(current_xlim, current_ylim))
+
+
     def highlight_value_in_tables(self, clicked_value):
         """Helper method to find and highlight a value in both tables."""
         
@@ -1622,6 +1976,11 @@ class ImageViewerWindow(QMainWindow):
                 
                 self.ax.set_xlim([xdata - x_range, xdata + x_range])
                 self.ax.set_ylim([ydata - y_range, ydata + y_range])
+
+                self.zoom_changed = True  # Flag that zoom has changed
+                
+                if not hasattr(self, 'zoom_changed'):
+                    self.zoom_changed = False
                 
             elif event.button == 3:  # Right click - zoom out
                 x_range = (current_xlim[1] - current_xlim[0])
@@ -1639,6 +1998,11 @@ class ImageViewerWindow(QMainWindow):
                 else:
                     self.ax.set_xlim(new_xlim)
                     self.ax.set_ylim(new_ylim)
+
+
+                self.zoom_changed = False  # Flag that zoom has changed
+                
+                
             
             self.canvas.draw()
         
@@ -1652,7 +2016,7 @@ class ImageViewerWindow(QMainWindow):
                     x_idx = int(round(event.xdata))
                     y_idx = int(round(event.ydata))
                     # Check if Ctrl key is pressed (using matplotlib's key_press system)
-                    ctrl_pressed = 'ctrl' in event.modifiers  # Note: changed from 'control' to 'ctrl'
+                    ctrl_pressed = 'ctrl' in event.modifiers  
                     if self.channel_data[self.active_channel][self.current_slice, y_idx, x_idx] != 0:
                         clicked_value = self.channel_data[self.active_channel][self.current_slice, y_idx, x_idx]
                     else:
@@ -2055,8 +2419,9 @@ class ImageViewerWindow(QMainWindow):
 
     def show_thresh_dialog(self):
         """Show threshold dialog"""
-        thresh_window = ThresholdWindow(self)
-        thresh_window.show()  # Non-modal window
+        dialog = ThresholdDialog(self)
+        dialog.exec()
+
 
     def show_mask_dialog(self):
         """Show the mask dialog"""
@@ -2265,13 +2630,10 @@ class ImageViewerWindow(QMainWindow):
 
 
                 except Exception as e:
-                    import traceback
-                    print(traceback.format_exc())
                     print(f"An error has occured: {e}")
 
             except Exception as e:
-                import traceback
-                print(traceback.format_exc())
+
                 QMessageBox.critical(
                     self,
                     "Error Loading",
@@ -2513,12 +2875,22 @@ class ImageViewerWindow(QMainWindow):
         """Shows a dialog asking user to confirm if image is 2D RGB"""
         msg = QMessageBox()
         msg.setIcon(QMessageBox.Icon.Question)
-        msg.setText("Image Format Detection")
+        msg.setText("Image Format Alert")
         msg.setInformativeText("Is this a 2D color (RGB/CMYK) image?")
         msg.setWindowTitle("Confirm Image Format")
         msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
         return msg.exec() == QMessageBox.StandardButton.Yes
 
+    def confirm_resize_dialog(self):
+        """Shows a dialog asking user to resize image"""
+        msg = QMessageBox()
+        msg.setIcon(QMessageBox.Icon.Question)
+        msg.setText("Image Format Alert")
+        msg.setInformativeText(f"This image is a different shape than the ones loaded into the viewer window. Trying to run processes with images of different sizes has a high probability of crashing the program.\nPress yes to resize the new image to the other images. Press no to load it anyway.")
+        msg.setWindowTitle("Resize")
+        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):
         """Load a channel and enable active channel selection if needed."""
 
@@ -2532,6 +2904,7 @@ class ImageViewerWindow(QMainWindow):
                     "TIFF Files (*.tif *.tiff)"
                 )
                 self.channel_data[channel_index] = tifffile.imread(filename)
+
                 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)
 
@@ -2547,6 +2920,23 @@ class ImageViewerWindow(QMainWindow):
             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])
 
+            reset_resize = False
+
+            for i in range(4): #Try to ensure users don't load in different sized arrays
+                if self.channel_data[i] is None or i == channel_index or data:
+                    if self.highlight_overlay is not None: #Make sure highlight overlay is always the same shape as new images
+                        if self.channel_data[i].shape[:3] != self.highlight_overlay.shape:
+                            self.resizing = True
+                            reset_resize = True
+                            self.highlight_overlay = None
+                    continue
+                else:
+                    old_shape = self.channel_data[i].shape[:3] #Ask user to resize images that are shaped differently
+                    if old_shape != self.channel_data[channel_index].shape[:3]:
+                        if self.confirm_resize_dialog():
+                            self.channel_data[channel_index] = n3d.upsample_with_padding(self.channel_data[channel_index], original_shape = old_shape)
+                            break
+
 
             if channel_index == 0:
                 my_network.nodes = self.channel_data[channel_index]
@@ -2601,11 +2991,14 @@ class ImageViewerWindow(QMainWindow):
             self.volume_dict[channel_index] = None #reset volumes
 
             if assign_shape: #keep original shape tracked to undo resampling.
-                self.original_shape = self.channel_data[channel_index].shape
+                if self.original_shape is None:
+                    self.original_shape = self.channel_data[channel_index].shape
+                elif self.original_shape[0] < self.channel_data[channel_index].shape[0] or self.original_shape[1] < self.channel_data[channel_index].shape[1] or self.original_shape[2] < self.channel_data[channel_index].shape[2]:
+                    self.original_shape = self.channel_data[channel_index].shape
                 if len(self.original_shape) == 4:
                     self.original_shape = (self.original_shape[0], self.original_shape[1], self.original_shape[2])
 
-            self.update_display()
+            self.update_display(reset_resize = reset_resize)
 
 
                 
@@ -2829,13 +3222,21 @@ class ImageViewerWindow(QMainWindow):
         current_ylim = self.ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
         # Convert slider values (0-100) to data values (0-1)
         min_val, max_val = values
-        self.channel_brightness[channel_index]['min'] = min_val / 255
+        self.channel_brightness[channel_index]['min'] = min_val / 255 #Accomodate 32 bit data?
         self.channel_brightness[channel_index]['max'] = max_val / 255
         self.update_display(preserve_zoom = (current_xlim, current_ylim))
+
+
+
     
-    def update_display(self, preserve_zoom=None, dims = None):
+    def update_display(self, preserve_zoom=None, dims = None, called = False, reset_resize = False, begin_paint = False):
         """Update the display with currently visible channels and highlight overlay."""
 
+        if begin_paint:
+            # Store/update the static background with current zoom level
+            self.static_background = self.canvas.copy_from_bbox(self.ax.bbox)
+
+
         self.figure.clear()
 
         # Get active channels and their dimensions
@@ -2927,6 +3328,9 @@ class ImageViewerWindow(QMainWindow):
                                  vmax=1,
                                  extent=(-0.5, min_width-0.5, min_height-0.5, -0.5))
 
+        if self.preview and not called:
+            self.create_highlight_overlay_slice(self.targs, bounds = self.bounds)
+
         # Add highlight overlay if it exists
         if self.highlight_overlay is not None:
             highlight_slice = self.highlight_overlay[self.current_slice]
@@ -2939,6 +3343,7 @@ class ImageViewerWindow(QMainWindow):
                          alpha=0.5)
 
 
+
         
         # Style the axes
         self.ax.set_xlabel('X')
@@ -2987,9 +3392,42 @@ class ImageViewerWindow(QMainWindow):
         if current_xlim is not None and current_ylim is not None:
             self.ax.set_xlim(current_xlim)
             self.ax.set_ylim(current_ylim)
+        if reset_resize:
+            self.resizing = False
         
         self.canvas.draw()
 
+    def update_display_slice(self, channel, preserve_zoom=None):
+        """Ultra minimal update that only changes the paint channel's data"""
+        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
+            channel_image.set_array(self.channel_data[channel][self.current_slice])
+            
+            # Restore the static background (all other channels) at current zoom level
+            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()
+
     def show_netshow_dialog(self):
         dialog = NetShowDialog(self)
         dialog.exec()
@@ -4162,17 +4600,18 @@ class WhiteDialog(QDialog):
     def white_overlay(self):
 
         try:
-
-            try:
+            if isinstance(my_network.nodes, np.ndarray) :
                 overlay = np.ones_like(my_network.nodes).astype(np.uint8) * 255
-            except:
+            elif isinstance(my_network.edges, np.ndarray): 
                 overlay = np.ones_like(my_network.edges).astype(np.uint8) * 255
-            finally:
-                my_network.id_overlay = overlay
+            elif isinstance(my_network.network_overlay, np.ndarray):
+                overlay = np.ones_like(my_network.network_overlay).astype(np.uint8) * 255
 
-                self.parent().load_channel(3, channel_data = my_network.id_overlay, data = True)
+            my_network.id_overlay = overlay
 
-                self.accept()
+            self.parent().load_channel(3, channel_data = my_network.id_overlay, data = True)
+
+            self.accept()
 
         except Exception as e:
             print(f"Error making white background: {e}")
@@ -4274,7 +4713,7 @@ class NetShowDialog(QDialog):
         self.geo_layout = QPushButton("geo_layout")
         self.geo_layout.setCheckable(True)
         self.geo_layout.setChecked(False)
-        layout.addRow("Use Geometric Layout:", self.geo_layout)
+        layout.addRow("Use Geographic Layout:", self.geo_layout)
         
         # Add mode selection dropdown
         self.mode_selector = QComboBox()
@@ -5246,101 +5685,504 @@ class LabelDialog(QDialog):
                 f"Error running label: {str(e)}"
             )
 
-class ThresholdWindow(QMainWindow):
+class ThresholdDialog(QDialog):
     def __init__(self, parent=None):
         super().__init__(parent)
-        self.setWindowTitle("Threshold Params (Active Image)")
-        
-        # Create central widget and layout
-        central_widget = QWidget()
-        self.setCentralWidget(central_widget)
-        layout = QFormLayout(central_widget)
-
-        self.min = QLineEdit("")
-        layout.addRow("Minimum Value to retain:", self.min)
+        self.setWindowTitle("Choose Threshold Mode")
+        self.setModal(True)
         
-        # Create widgets
-        self.max = QLineEdit("")
-        layout.addRow("Maximum Value to retain:", self.max)
+        layout = QFormLayout(self)
 
         # Add mode selection dropdown
         self.mode_selector = QComboBox()
-        self.mode_selector.addItems(["Using Volumes", "Using Label/Brightness"])
+        self.mode_selector.addItems(["Using Label/Brightness", "Using Volumes"])
         self.mode_selector.setCurrentIndex(0)  # Default to Mode 1
         layout.addRow("Execution Mode:", self.mode_selector)
 
         # Add Run button
-        prev_button = QPushButton("Preview")
-        prev_button.clicked.connect(self.run_preview)
-        layout.addRow(prev_button)
+        run_button = QPushButton("Select")
+        run_button.clicked.connect(self.thresh_mode)
+        layout.addRow(run_button)
+
+        # Add ML button
+        ML = QPushButton("Machine Learning")
+        ML.clicked.connect(self.start_ml)
+        layout.addRow(ML)
+
+
+    def thresh_mode(self):
+
+        try:
+
+            accepted_mode = self.mode_selector.currentIndex()
+
+            if accepted_mode == 1:
+                if len(np.unique(self.parent().channel_data[self.parent().active_channel])) < 3:
+                    self.parent().show_label_dialog()
+
+                if self.parent().volume_dict[self.parent().active_channel] is None:
+                    self.parent().volumes()
+
+            thresh_window = ThresholdWindow(self.parent(), accepted_mode)
+            thresh_window.show()  # Non-modal window
+            self.highlight_overlay = None
+            self.accept()
+        except:
+            pass
+
+    def start_ml(self):
+
+
+        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():
+                pass
+            else:
+                return
+        elif self.parent().channel_data[0] is None and self.parent().channel_data[1] is None:
+            QMessageBox.critical(
+                self,
+                "Alert",
+                "Requires the channel for segmentation to be loaded into either the nodes or edges channels"
+            )
+            return
+
+
+        self.parent().machine_window = MachineWindow(self.parent())
+        self.parent().machine_window.show()  # Non-modal window
+        self.accept()
+
+    def confirm_machine_dialog(self):
+        """Shows a dialog asking user to confirm if they want to start the segmenter"""
+        msg = QMessageBox()
+        msg.setIcon(QMessageBox.Icon.Question)
+        msg.setText("Alert")
+        msg.setInformativeText("Use of this feature will require use of both overlay channels and the highlight overlay. Please save any data and return, or proceed if you do not need those overlays")
+        msg.setWindowTitle("Proceed?")
+        msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
+        return msg.exec() == QMessageBox.StandardButton.Yes
+
+
+class MachineWindow(QMainWindow):
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+
+        self.setWindowTitle("Threshold")
         
-        # Add Run button
+        # 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)
+
+
+
+
+        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]
+
+
+        array1 = np.zeros_like(active_data)
+        array2 = np.zeros_like(active_data)
+        array3 = np.zeros_like(active_data)
+        self.parent().highlight_overlay = array3 #Clear this out for the segmenter to use
+
+        self.parent().load_channel(2, array1, True) #Temp for debugging
+        self.parent().load_channel(3, array2, True) #Temp for debugging
+
+        self.parent().base_colors[2] = self.parent().color_dictionary['LIGHT_GREEN']
+        self.parent().base_colors[3] = self.parent().color_dictionary['LIGHT_RED']
+
+        
+        # Set a reasonable default size
+        self.setMinimumWidth(400)
+        self.setMinimumHeight(400)
+
+        # Create zoom button and pan button
+        buttons_widget = QWidget()
+        buttons_layout = QHBoxLayout(buttons_widget)
+
+        # Create zoom 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)
+        form_layout.addWidget(self.brush_button)
+        self.brush_button.click()
+
+        self.fore_button = QPushButton("Foreground")
+        self.fore_button.setCheckable(True)
+        self.fore_button.setChecked(True)
+        self.fore_button.clicked.connect(self.toggle_foreground)
+        form_layout.addWidget(self.fore_button)
+
+        self.back_button = QPushButton("Background")
+        self.back_button.setCheckable(True)
+        self.back_button.setChecked(False)
+        self.back_button.clicked.connect(self.toggle_background)
+        form_layout.addWidget(self.back_button)
+
+        train_button = QPushButton("Train Model")
+        train_button.clicked.connect(self.train_model)
+        form_layout.addRow(train_button)
+
+        seg_button = QPushButton("Segment")
+        seg_button.clicked.connect(self.segment)
+        form_layout.addRow(seg_button)
+
+        self.trained = False
+
+
+        self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=True)
+
+
+
+
+    def toggle_foreground(self):
+
+        self.parent().foreground = self.fore_button.isChecked()
+
+        if self.parent().foreground:
+            self.back_button.setChecked(False)
+        else:
+            self.back_button.setChecked(True)
+
+    def switch_foreground(self):
+
+        self.fore_button.click()
+
+    def toggle_background(self):
+
+        self.parent().foreground = not self.back_button.isChecked()
+
+        if not self.parent().foreground:
+            self.fore_button.setChecked(False)
+        else:
+            self.fore_button.setChecked(True)
+
+
+
+    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().pan_button.setChecked(False)
+            self.parent().zoom_button.setChecked(False)
+            self.parent().pan_mode = False
+            self.parent().zoom_mode = False
+            self.parent().update_brush_cursor()
+        else:
+            self.parent().zoom_button.click()
+
+    def silence_button(self):
+        self.brush_button.setChecked(False)
+
+    def toggle_brush_button(self):
+
+        self.brush_button.click()
+
+    def train_model(self):
+
+        self.segmenter.train_batch(self.parent().channel_data[2], self.parent().channel_data[3])
+        self.trained = True
+
+    def segment(self):
+
+        if not self.trained:
+            return
+        else:
+            foreground_coords, background_coords = self.segmenter.segment_volume()
+
+            # Clean up when done
+            self.segmenter.cleanup()
+
+        for z,y,x in foreground_coords:
+            self.parent().highlight_overlay[z,y,x] = True
+
+        self.parent().update_display()
+
+    def closeEvent(self, event):
+        if self.brush_button.isChecked():
+            self.silence_button()
+            self.toggle_brush_mode()
+        self.parent().brush_mode = False
+        self.parent().machine_window = None
+
+
+
+
+
+class ThresholdWindow(QMainWindow):
+    def __init__(self, parent=None, accepted_mode=0):
+        super().__init__(parent)
+        self.setWindowTitle("Threshold")
+        
+        # Create central widget and layout
+        central_widget = QWidget()
+        self.setCentralWidget(central_widget)
+        layout = QVBoxLayout(central_widget)
+        
+        # Get histogram data
+        if accepted_mode == 1:
+            self.histo_list = list(self.parent().volume_dict[self.parent().active_channel].values())
+            self.bounds = False
+            self.parent().bounds = False
+        elif accepted_mode == 0:
+            targ_shape = self.parent().channel_data[self.parent().active_channel].shape
+            if (targ_shape[0] + targ_shape[1] + targ_shape[2]) > 2500: #Take a simpler histogram on big arrays
+                temp_max = np.max(self.parent().channel_data[self.parent().active_channel])
+                temp_min = np.min(self.parent().channel_data[self.parent().active_channel])
+                temp_array = n3d.downsample(self.parent().channel_data[self.parent().active_channel], 5)
+                self.histo_list = temp_array.flatten().tolist()
+                self.histo_list.append(temp_min)
+                self.histo_list.append(temp_max)
+            else: #Otherwise just use full array data
+                self.histo_list = self.parent().channel_data[self.parent().active_channel].flatten().tolist()
+            self.bounds = True
+            self.parent().bounds = True
+            
+        # Create matplotlib figure
+        fig = Figure(figsize=(5, 4))
+        self.canvas = FigureCanvas(fig)
+        layout.addWidget(self.canvas)
+        
+        # Pre-compute histogram with numpy
+        counts, bin_edges = np.histogram(self.histo_list, bins=50)
+        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
+        
+        # Plot pre-computed histogram
+        self.ax = fig.add_subplot(111)
+        self.ax.bar(bin_centers, counts, width=bin_edges[1] - bin_edges[0], alpha=0.5)
+        
+        # Add vertical lines for thresholds
+        self.min_line = self.ax.axvline(min(self.histo_list), color='r')
+        self.max_line = self.ax.axvline(max(self.histo_list), color='b')
+        
+        # Connect events for dragging
+        self.canvas.mpl_connect('button_press_event', self.on_press)
+        self.canvas.mpl_connect('motion_notify_event', self.on_motion)
+        self.canvas.mpl_connect('button_release_event', self.on_release)
+        
+        self.dragging = None
+        
+        # Store histogram bounds
+        if self.bounds:
+            self.data_min = 0
+        else:
+            self.data_min = min(self.histo_list)
+        self.data_max = max(self.histo_list)
+
+        # Create form layout for inputs
+        form_layout = QFormLayout()
+        
+        self.min = QLineEdit(f"{self.data_min}")
+        self.min.editingFinished.connect(self.min_value_changed)
+        form_layout.addRow("Minimum Value to retain:", self.min)
+        self.prev_min = self.data_min
+        
+        self.max = QLineEdit(f"{self.data_max}")
+        self.max.editingFinished.connect(self.max_value_changed)
+        form_layout.addRow("Maximum Value to retain:", self.max)
+        self.prev_max = self.data_max
+
+        self.targs = [self.prev_min, self.prev_max]
+        
+        # preview checkbox (default False)
+        self.preview = QPushButton("Preview")
+        self.preview.setCheckable(True)
+        self.preview.setChecked(False)
+        self.preview.clicked.connect(self.preview_mode)
+        form_layout.addRow("Show Preview:", self.preview)
+
         run_button = QPushButton("Apply Threshold")
         run_button.clicked.connect(self.thresh)
-        layout.addRow(run_button)
+        form_layout.addRow(run_button)
         
-        # Set a reasonable default size
-        self.setMinimumWidth(300)
+        layout.addLayout(form_layout)
         
-    def run_preview(self):
+        # Set a reasonable default size
+        self.setMinimumWidth(400)
+        self.setMinimumHeight(400)
 
-        def get_valid_float(text, default_value):
-            try:
-                return float(text) if text.strip() else default_value
-            except ValueError:
-                print(f"Invalid input: {text}")
-                return default_value
+    def closeEvent(self, event):
+        self.parent().preview = False
+        self.parent().targs = None
+        self.parent().bounds = False
 
+    def get_values_in_range(self, lst, min_val, max_val):
+        values = [x for x in lst if min_val <= x <= max_val]
+        output = []
+        for item in self.parent().volume_dict[self.parent().active_channel]:
+            if self.parent().volume_dict[self.parent().active_channel][item] in values:
+                output.append(item)
+        return output
+
+
+    def min_value_changed(self):
         try:
-            channel = self.parent().active_channel
-            accepted_mode = self.mode_selector.currentIndex()
+            text = self.min.text()
+            if not text:  # If empty, ignore
+                return
             
-            if accepted_mode == 0:
-                if len(np.unique(self.parent().channel_data[self.parent().active_channel])) < 3:
-                    self.parent().show_label_dialog()
+            try:
+                value = float(text)
+                
+                # Bound check against data limits
+                value = max(self.data_min, value)
+
+                # Check against max line
+                max_val = float(self.max.text()) if self.max.text() else self.data_max
+                if value > max_val:
+                    # If min would exceed max, set max to its highest possible value
+                    self.max.setText(str(round(self.data_max, 2)))
+                    self.max_line.set_xdata([self.data_max, self.data_max])
+                    # And set min to the previous max value
+                    value = max_val
+                    self.min.setText(str(round(value, 2)))
+
+                if value == self.prev_min:
+                    return
+                else:
+                    self.prev_min = value
+                    if self.bounds:
+                        self.targs = [self.prev_min, self.prev_max]
+                    else:
+                        self.targs = self.get_values_in_range(self.histo_list, self.prev_min, self.prev_max)
+                    self.parent().targs = self.targs
+                    if self.preview.isChecked():
+                        self.parent().highlight_overlay = None
+                        self.parent().create_highlight_overlay_slice(self.targs, bounds = self.bounds)
+                
+                # Update the line
+                self.min_line.set_xdata([value, value])
+                self.canvas.draw()
+
 
-                if self.parent().volume_dict[channel] is None:
-                    self.parent().volumes()
-                    
-                volumes = self.parent().volume_dict[channel]
-                default_max = max(volumes.values())
-                default_min = min(volumes.values())
                 
-                max_val = get_valid_float(self.max.text(), default_max)
-                min_val = get_valid_float(self.min.text(), default_min)
+            except ValueError:
+                # If invalid number, reset to current line position
+                self.min.setText(str(round(self.min_line.get_xself.data_mindata()[0], 2)))
+        except:
+            pass
+
+    def max_value_changed(self):
+        try:
+            text = self.max.text()
+            if not text:  # If empty, ignore
+                return
                 
-                valid_indices = [item for item in volumes 
-                                if min_val <= volumes[item] <= max_val]
-                                
-            elif accepted_mode == 1:
-                channel_data = self.parent().channel_data[self.parent().active_channel]
-                default_max = np.max(channel_data)
-                default_min = np.min(channel_data)
+            try:
+                value = float(text)
                 
-                max_val = int(get_valid_float(self.max.text(), default_max))
-                min_val = int(get_valid_float(self.min.text(), default_min))
+                # Bound check against data limits
+                value = min(self.data_max, value)
                 
-                if min_val > max_val:
-                    min_val, max_val = max_val, min_val
-                    
-                valid_indices = list(range(min_val, max_val + 1))
+                # Check against min line
+                min_val = float(self.min.text()) if self.min.text() else self.data_min
+                if value < min_val:
+                    # If max would go below min, set min to its lowest possible value
+                    self.min.setText(str(round(self.data_min, 2)))
+                    self.min_line.set_xdata([self.data_min, self.data_min])
+                    # And set max to the previous min value
+                    value = min_val
+                    self.max.setText(str(round(value, 2)))
+
+                if value == self.prev_max:
+                    return
+                else:
+                    self.prev_max = value
+                    if self.bounds:
+                        self.targs = [self.prev_min, self.prev_max]
+                    else:
+                        self.targs = self.get_values_in_range(self.histo_list, self.prev_min, self.prev_max)
+                    self.parent().targs = self.targs
+                    if self.preview.isChecked():
+                        self.parent().highlight_overlay = None
+                        self.parent().create_highlight_overlay_slice(self.targs, bounds = self.bounds)
+                
+                # Update the line
+                self.max_line.set_xdata([value, value])
+                self.canvas.draw()
+
+
+
+
+
+                
+            except ValueError:
+                # If invalid number, reset to current line position
+                self.max.setText(str(round(self.max_line.get_xdata()[0], 2)))
+        except:
+            pass
+        
+    def on_press(self, event):
+        try:
+            if event.inaxes != self.ax:
+                return
             
-            if channel == 0:
-                self.parent().create_highlight_overlay(node_indices = valid_indices)
-            elif channel == 1:
-                self.parent().create_highlight_overlay(edge_indices = valid_indices)
-            elif channel == 2:
-                self.parent().create_highlight_overlay(overlay1_indices = valid_indices)
-            elif channel == 3:
-                self.parent().create_highlight_overlay(overlay2_indices = valid_indices)
+            # Left click controls left line
+            if event.button == 1:  # Left click
+                self.dragging = 'min'
+            # Right click controls right line
+            elif event.button == 3:  # Right click
+                self.dragging = 'max'
+        except:
+            pass
+                
+    def on_motion(self, event):
+        try:
+            if not self.dragging or event.inaxes != self.ax:
+                return
+                
+            if self.dragging == 'min':
+                if event.xdata < self.max_line.get_xdata()[0]:
+                    self.min_line.set_xdata([event.xdata, event.xdata])
+                    self.min.setText(str(round(event.xdata, 2)))
+            else:
+                if event.xdata > self.min_line.get_xdata()[0]:
+                    self.max_line.set_xdata([event.xdata, event.xdata])
+                    self.max.setText(str(round(event.xdata, 2)))
+                    
+            self.canvas.draw()
+        except:
+            pass
+        
+    def on_release(self, event):
+        self.min_value_changed()
+        self.max_value_changed()
+        self.dragging = None
 
-        except Exception as e:
-            print(f"Error showing preview: {e}")
+    def preview_mode(self):
+        try:
+            preview = self.preview.isChecked()
+            self.parent().preview = preview
+            self.parent().targs = self.targs
+
+            if preview and self.targs is not None:
+                self.parent().create_highlight_overlay_slice(self.parent().targs, bounds = self.bounds)
+        except:
+            pass      
 
     def thresh(self):
         try:
 
-            self.run_preview()
+            if self.parent().active_channel == 0:
+                self.parent().create_highlight_overlay(node_indices = self.targs, bounds = self.bounds)
+            elif self.parent().active_channel == 1:
+                self.parent().create_highlight_overlay(edge_indices = self.targs, bounds = self.bounds)
+            elif self.parent().active_channel == 2:
+                self.parent().create_highlight_overlay(overlay1_indices = self.targs, bounds = self.bounds)
+            elif self.parent().active_channel == 3:
+                self.parent().create_highlight_overlay(overlay2_indices = self.targs, bounds = self.bounds)
+
             channel_data = self.parent().channel_data[self.parent().active_channel]
             mask = self.parent().highlight_overlay > 0
             channel_data = channel_data * mask
@@ -5349,6 +6191,7 @@ class ThresholdWindow(QMainWindow):
             self.close()
             
         except Exception as e:
+
             QMessageBox.critical(
                 self,
                 "Error",
@@ -5775,9 +6618,17 @@ class WatershedDialog(QDialog):
         self.directory = QLineEdit()
         self.directory.setPlaceholderText("Leave empty for None")
         layout.addRow("Output Directory:", self.directory)
+
+        active_shape = self.parent().channel_data[self.parent().active_channel].shape[0]
+
+        if active_shape == 1:
+            self.default = 0.2
+        else:
+            self.default = 0.05
+
         
         # Proportion (default 0.1)
-        self.proportion = QLineEdit("0.05")
+        self.proportion = QLineEdit(f"{self.default}")
         layout.addRow("Proportion:", self.proportion)
         
         # GPU checkbox (default True)
@@ -5813,9 +6664,9 @@ class WatershedDialog(QDialog):
             
             # Get proportion (0.1 if empty or invalid)
             try:
-                proportion = float(self.proportion.text()) if self.proportion.text() else 0.05
+                proportion = float(self.proportion.text()) if self.proportion.text() else self.default
             except ValueError:
-                proportion = 0.05
+                proportion = self.default
             
             # Get GPU state
             gpu = self.gpu.isChecked()
@@ -5842,7 +6693,8 @@ class WatershedDialog(QDialog):
             active_data = self.parent().channel_data[self.parent().active_channel]
             if active_data is None:
                 raise ValueError("No active image selected")
-            
+
+
             # Call watershed method with parameters
             result = n3d.watershed(
                 active_data,
@@ -5865,6 +6717,8 @@ class WatershedDialog(QDialog):
             self.accept()
             
         except Exception as e:
+            import traceback
+            print(traceback.format_exc())
             QMessageBox.critical(
                 self,
                 "Error",
@@ -6119,11 +6973,21 @@ class BranchDialog(QDialog):
         self.nodes.setChecked(True)
         layout.addRow("Generate nodes from edges? (Skip if already completed - presumes your edge skeleton from generate nodes is in Edges and that your original Edges are in Overlay 2):", self.nodes)
 
-        # GPU checkbox (default True)
+        # GPU checkbox (default False)
         self.GPU = QPushButton("GPU")
         self.GPU.setCheckable(True)
         self.GPU.setChecked(False)
-        layout.addRow("Use GPU (Note this may need to temporarily downsample your large images which may simplify outputs - Only memory errors but not permission errors for accessing GRAM are handled by default - CPU will never try to downsample):", self.GPU)
+        layout.addRow("Use GPU (Note this may need to temporarily downsample your large images which may simplify outputs - Only memory errors but not permission errors for accessing VRAM are handled by default - CPU will never try to downsample):", self.GPU)
+
+        # Branch Fix checkbox (default False)
+        self.fix = QPushButton("Auto-Correct Branches")
+        self.fix.setCheckable(True)
+        self.fix.setChecked(False)
+        layout.addRow("Attempt to auto-correct branch labels:", self.fix)
+
+        self.fix_val = QLineEdit()
+        self.fix_val.setPlaceholderText("Empty = default value...")
+        layout.addRow("If checked above - Avg Degree of Nearby Branch Communities to Merge (Attempt to fix branch labeling - try 4 to 6 to start or leave empty):", self.fix_val)
 
         self.down_factor = QLineEdit("0")
         layout.addRow("Internal downsample (will have to recompute nodes)?:", self.down_factor)
@@ -6151,6 +7015,9 @@ class BranchDialog(QDialog):
             nodes = self.nodes.isChecked()
             GPU = self.GPU.isChecked()
             cubic = self.cubic.isChecked()
+            fix = self.fix.isChecked()
+            fix_val = float(self.fix_val.text()) if self.fix_val.text() else None
+
 
 
             original_shape = my_network.edges.shape
@@ -6158,7 +7025,7 @@ class BranchDialog(QDialog):
 
             if down_factor > 0:
                 self.parent().show_gennodes_dialog(down_factor = [down_factor, cubic], called = True)
-            elif nodes:
+            elif nodes or my_network.nodes is None:
                 self.parent().show_gennodes_dialog(called = True)
                 down_factor = None
 
@@ -6166,6 +7033,21 @@ class BranchDialog(QDialog):
 
                 output = n3d.label_branches(my_network.edges, nodes = my_network.nodes, bonus_array = original_array, GPU = GPU, down_factor = down_factor, arrayshape = original_shape)
 
+                if fix:
+
+                    temp_network = n3d.Network_3D(nodes = output)
+
+                    temp_network.morph_proximity(search = 1) #Detect network of nearby branches
+
+                    temp_network.community_partition(weighted = False, style = 1, dostats = False) #Find communities with louvain, unweighted params
+
+                    targs = n3d.fix_branches(temp_network.nodes, temp_network.network, temp_network.communities, fix_val)
+
+                    temp_network.com_to_node(targs)
+
+                    output = temp_network.nodes
+
+
                 if down_factor is not None:
 
                     self.parent().reset(nodes = True, id_overlay = True, edges = True)
@@ -6181,6 +7063,8 @@ class BranchDialog(QDialog):
 
         except Exception as e:
             print(f"Error labeling branches: {e}")
+            import traceback
+            print(traceback.format_exc())
 
 
 
@@ -6336,7 +7220,7 @@ class AlterDialog(QDialog):
 class ModifyDialog(QDialog):
     def __init__(self, parent=None):
         super().__init__(parent)
-        self.setWindowTitle("Create Nodes from Edge Vertices")
+        self.setWindowTitle("Modify Network Qualities")
         self.setModal(True)
         layout = QFormLayout(self)
         
@@ -6376,6 +7260,12 @@ class ModifyDialog(QDialog):
         self.isolate.setChecked(False)
         layout.addRow("Isolate connections between two specific node types (if assigned)?:", self.isolate)
 
+        # Community collapse checkbox (default False)
+        self.comcollapse = QPushButton("Communities -> nodes")
+        self.comcollapse.setCheckable(True)
+        self.comcollapse.setChecked(False)
+        layout.addRow("Convert communities to nodes?:", self.comcollapse)
+
         #change button
         change_button = QPushButton("Add/Remove Network Pairs")
         change_button.clicked.connect(self.show_alter_dialog)
@@ -6409,6 +7299,8 @@ class ModifyDialog(QDialog):
             edgeweight = self.edgeweight.isChecked()
             prune = self.prune.isChecked()
             isolate = self.isolate.isChecked()
+            comcollapse = self.comcollapse.isChecked()
+
 
             if isolate and my_network.node_identities is not None:
                 self.show_isolate_dialog()
@@ -6434,6 +7326,15 @@ class ModifyDialog(QDialog):
                     self.parent().format_for_upperright_table(my_network.node_centroids, 'NodeID', ['Z', 'Y', 'X'], 'Node Centroids')
                 except:
                     pass
+            if comcollapse:
+                if my_network.communities is None:
+                    self.parent().show_partition_dialog()
+                    if my_network.communities is None:
+                        return
+                my_network.com_to_node()
+                self.parent().load_channel(0, my_network.nodes, True)
+                my_network.communities = None
+
             try:
                 if hasattr(my_network, 'network_lists'):
                     model = PandasModel(my_network.network_lists)
@@ -6552,8 +7453,7 @@ class CentroidDialog(QDialog):
                 "Error",
                 f"Error finding centroids: {str(e)}"
             )
-            import traceback
-            print(traceback.format_exc())
+