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

nettracer3d/nettracer.py

@@ -858,10 +858,26 @@ def hash_inners(search_region, inner_edges, GPU = True):
 
     return inner_edges
 
+
+def dilate_2D(array, search, scaling = 1):
+
+    inv = array < 1
+
+    inv = smart_dilate.compute_distance_transform_distance(inv)
+
+    inv = inv * scaling
+
+    inv = inv <= search
+
+    return inv
+
+
 def dilate_3D(tiff_array, dilated_x, dilated_y, dilated_z):
     """Internal method to dilate an array in 3D. Dilation this way is much faster than using a distance transform although the latter is theoretically more accurate.
     Arguments are an array,  and the desired pixel dilation amounts in X, Y, Z."""
 
+    if tiff_array.shape[0] == 1:
+        return dilate_2D(tiff_array, ((dilated_x - 1) / 2))
 
     if dilated_x == 3 and dilated_y == 3  and dilated_z == 3:
 
@@ -996,6 +1012,8 @@ def dilate_3D_recursive(tiff_array, dilated_x, dilated_y, dilated_z, step_size=N
     
     Each dilation parameter represents (n-1)/2 steps outward from the object.
     """
+    if tiff_array.shape[0] == 1:
+        return dilate_2D(tiff_array, ((dilated_x - 1) / 2))
     # Calculate the smallest dimension of the array
     min_dim = min(tiff_array.shape)
     
@@ -2414,6 +2432,7 @@ class Network_3D:
         Can be called on a Network_3D object to save the node centroids properties to hard mem as a .xlsx file. It will save to the active directory if none is specified.
         :param directory: (Optional - Val = None; String). The path to an indended directory to save the centroids to.
         """
+
         if self._node_centroids is not None:
             if directory is None:
                 network_analysis._save_centroid_dictionary(self._node_centroids, 'node_centroids.xlsx')
@@ -2424,7 +2443,14 @@ class Network_3D:
                 print(f"Centroids saved to {directory}/node_centroids.xlsx")
 
         if self._node_centroids is None:
-            print("Node centroids attribute is empty, did not save...")
+            if directory is None:
+                network_analysis._save_centroid_dictionary({}, 'node_centroids.xlsx')
+                print("Centroids saved to node_centroids.xlsx")
+
+            if directory is not None:
+                network_analysis._save_centroid_dictionary({}, f'{directory}/node_centroids.xlsx')
+                print(f"Centroids saved to {directory}/node_centroids.xlsx")
+
 
     def save_edge_centroids(self, directory = None):
         """
@@ -2442,6 +2468,13 @@ class Network_3D:
 
         if self._edge_centroids is None:
             print("Edge centroids attribute is empty, did not save...")
+            if directory is None:
+                network_analysis._save_centroid_dictionary({}, 'edge_centroids.xlsx', index = 'Edge ID')
+                print("Centroids saved to edge_centroids.xlsx")
+
+            if directory is not None:
+                network_analysis._save_centroid_dictionary({}, f'{directory}/edge_centroids.xlsx', index = 'Edge ID')
+                print(f"Centroids saved to {directory}/edge_centroids.xlsx")
 
     def save_search_region(self, directory = None):
         """
@@ -2494,7 +2527,13 @@ class Network_3D:
                 print(f"Node identities saved to {directory}/node_identities.xlsx")
 
         if self._node_identities is None:
-            print("Node identities attribute is empty...")
+            if directory is None:
+                network_analysis.save_singval_dict({}, 'NodeID', 'Identity', 'node_identities.xlsx')
+                print("Node identities saved to node_identities.xlsx")
+
+            if directory is not None:
+                network_analysis.save_singval_dict({}, 'NodeID', 'Identity', f'{directory}/node_identities.xlsx')
+                print(f"Node identities saved to {directory}/node_identities.xlsx")
 
     def save_communities(self, directory = None):
         """
@@ -2511,8 +2550,13 @@ class Network_3D:
                 print(f"Communities saved to {directory}/node_communities.xlsx")
 
         if self._communities is None:
-            print("Communities attribute is empty...")
+            if directory is None:
+                network_analysis.save_singval_dict({}, 'NodeID', 'Community', 'node_communities.xlsx')
+                print("Communities saved to node_communities.xlsx")
 
+            if directory is not None:
+                network_analysis.save_singval_dict({}, 'NodeID', 'Community', f'{directory}/node_communities.xlsx')
+                print(f"Communities saved to {directory}/node_communities.xlsx")
 
     def save_network_overlay(self, directory = None, filename = None):
 

nettracer3d/nettracer_gui.py

@@ -5,7 +5,7 @@ from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QFormLayout, QLineEdit, QPushButton, QFileDialog,
                             QLabel, QComboBox, QMessageBox, QTableView, QInputDialog,
                             QMenu, QTabWidget)
-from PyQt6.QtCore import (QPoint, Qt, QAbstractTableModel, QTimer)
+from PyQt6.QtCore import (QPoint, Qt, QAbstractTableModel, QTimer,  QThread, pyqtSignal)
 import numpy as np
 import time
 from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
@@ -13,8 +13,8 @@ from matplotlib.figure import Figure
 import matplotlib.pyplot as plt
 from qtrangeslider import QRangeSlider
 from nettracer3d import nettracer as n3d
-from nettracer3d import smart_dilate as sdl
 from nettracer3d import proximity as pxt
+from nettracer3d import smart_dilate as sdl
 from matplotlib.colors import LinearSegmentedColormap
 from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
 import pandas as pd
@@ -27,6 +27,7 @@ from functools import partial
 from nettracer3d import segmenter
 
 
+
 class ImageViewerWindow(QMainWindow):
     def __init__(self):
         super().__init__()
@@ -178,7 +179,7 @@ class ImageViewerWindow(QMainWindow):
         self.zoom_button.setCheckable(True)
         self.zoom_button.setFixedSize(40, 40)
         self.zoom_button.clicked.connect(self.toggle_zoom_mode)
-        control_layout.addWidget(self.zoom_button)
+        buttons_layout.addWidget(self.zoom_button)
         self.resizing = False
 
         self.pan_button = QPushButton("✋")
@@ -187,6 +188,14 @@ class ImageViewerWindow(QMainWindow):
         self.pan_button.clicked.connect(self.toggle_pan_mode)
         buttons_layout.addWidget(self.pan_button)
 
+        self.high_button = QPushButton("👁️")
+        self.high_button.setCheckable(True)
+        self.high_button.setFixedSize(40, 40)
+        self.high_button.clicked.connect(self.toggle_highlight)
+        self.high_button.setChecked(True)
+        buttons_layout.addWidget(self.high_button)
+        self.highlight = True
+
         control_layout.addWidget(buttons_widget)
 
         self.preview = False #Whether in preview mode or not
@@ -1429,6 +1438,12 @@ class ImageViewerWindow(QMainWindow):
 
 
 
+    def toggle_highlight(self):
+        self.highlight = self.high_button.isChecked()
+        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 toggle_zoom_mode(self):
@@ -1512,6 +1527,8 @@ class ImageViewerWindow(QMainWindow):
         if self.machine_window is not None:
             if event.key() == Qt.Key_A:
                 self.machine_window.switch_foreground()
+        if event.key() == Qt.Key_X:
+            self.high_button.click()
 
 
     def update_brush_cursor(self):
@@ -1649,11 +1666,7 @@ class ImageViewerWindow(QMainWindow):
                 x, y = int(event.xdata), int(event.ydata)
                 self.last_paint_pos = (x, y)
 
-                if self.foreground:
-                    channel = 2
-                else:
-                    channel = 3
-                
+                channel = 2
 
                 # Paint at initial position
                 self.paint_at_position(x, y, self.erase, channel)
@@ -1692,8 +1705,10 @@ class ImageViewerWindow(QMainWindow):
             
         if erase:
             val = 0
+        elif self.foreground:
+            val = 1
         else:
-            val = 255
+            val = 2
 
         height, width = self.channel_data[channel][self.current_slice].shape
         radius = self.brush_size // 2
@@ -1774,10 +1789,7 @@ class ImageViewerWindow(QMainWindow):
                 
             x, y = int(event.xdata), int(event.ydata)
 
-            if self.foreground:
-                channel = 2
-            else:
-                channel = 3
+            channel = 2
 
             
             if self.channel_data[2] is not None:
@@ -3313,26 +3325,41 @@ class ImageViewerWindow(QMainWindow):
                     else:
                         normalized_image = np.clip((current_image - vmin) / (vmax - vmin), 0, 1)
                     
-                    # Create custom colormap with higher intensity
-                    color = base_colors[channel]
-                    custom_cmap = LinearSegmentedColormap.from_list(
-                        f'custom_{channel}',
-                        [(0,0,0,0), (*color,1)]
-                    )
-                    
-                    # Display the image with slightly higher alpha
-                    self.ax.imshow(normalized_image,
-                                 alpha=0.7,
-                                 cmap=custom_cmap,
-                                 vmin=0,
-                                 vmax=1,
-                                 extent=(-0.5, min_width-0.5, min_height-0.5, -0.5))
+                    if channel == 2 and self.machine_window is not None:
+                        custom_cmap = LinearSegmentedColormap.from_list(
+                            f'custom_{channel}',
+                            [(0, 0, 0, 0),          # transparent for 0
+                             (0.5, 1, 0.5, 1),      # light green for 1
+                             (1, 0.5, 0.5, 1)]      # light red for 2
+                        )
+                        self.ax.imshow(current_image,
+                                     cmap=custom_cmap,
+                                     vmin=0,
+                                     vmax=2,
+                                     alpha=0.7,
+                                     interpolation='nearest',
+                                     extent=(-0.5, min_width-0.5, min_height-0.5, -0.5))
+                    else:
+                        # Create custom colormap with higher intensity
+                        color = base_colors[channel]
+                        custom_cmap = LinearSegmentedColormap.from_list(
+                            f'custom_{channel}',
+                            [(0,0,0,0), (*color,1)]
+                        )
+                        
+                        # Display the image with slightly higher alpha
+                        self.ax.imshow(normalized_image,
+                                     alpha=0.7,
+                                     cmap=custom_cmap,
+                                     vmin=0,
+                                     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:
+        if self.highlight_overlay is not None and self.highlight and self.machine_window is None:
             highlight_slice = self.highlight_overlay[self.current_slice]
             highlight_cmap = LinearSegmentedColormap.from_list(
                 'highlight',
@@ -3341,6 +3368,19 @@ class ImageViewerWindow(QMainWindow):
             self.ax.imshow(highlight_slice,
                          cmap=highlight_cmap,
                          alpha=0.5)
+        elif self.highlight_overlay is not None and self.highlight:
+            highlight_slice = self.highlight_overlay[self.current_slice]
+            highlight_cmap = LinearSegmentedColormap.from_list(
+                'highlight',
+                [(0, 0, 0, 0),          # transparent for 0
+                 (1, 1, 0, 1),          # bright yellow for 1
+                 (0, 0.7, 1, 1)]        # cool blue for 2
+            )
+            self.ax.imshow(highlight_slice,
+                         cmap=highlight_cmap,
+                         vmin=0,
+                         vmax=2,         # Important: set vmax to 2 to accommodate both values
+                         alpha=0.5)
 
 
 
@@ -5781,13 +5821,13 @@ class MachineWindow(QMainWindow):
         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]
+                act_channel = 0
             else:
                 active_data = self.parent().channel_data[1]
+                act_channel = 1
 
 
         array1 = np.zeros_like(active_data)
@@ -5795,12 +5835,18 @@ class MachineWindow(QMainWindow):
         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().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().base_colors[act_channel] = self.parent().color_dictionary['WHITE']
         self.parent().base_colors[2] = self.parent().color_dictionary['LIGHT_GREEN']
-        self.parent().base_colors[3] = self.parent().color_dictionary['LIGHT_RED']
 
+        self.parent().update_display()
         
         # Set a reasonable default size
         self.setMinimumWidth(400)
@@ -5830,20 +5876,39 @@ class MachineWindow(QMainWindow):
         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)
+        self.GPU = QPushButton("GPU")
+        self.GPU.setCheckable(True)
+        self.GPU.setChecked(False)
+        self.GPU.clicked.connect(self.toggle_GPU)
+        form_layout.addWidget(self.GPU)
+        self.use_gpu = False
+
+        train_button = QPushButton("Train Quick Model")
+        train_button.clicked.connect(lambda: self.train_model(speed = True))
         form_layout.addRow(train_button)
 
-        seg_button = QPushButton("Segment")
-        seg_button.clicked.connect(self.segment)
+        train_button = QPushButton("Train More Detailed Model")
+        train_button.clicked.connect(lambda: self.train_model(speed = False))
+        form_layout.addRow(train_button)
+
+        seg_button = QPushButton("Preview Segment")
+        seg_button.clicked.connect(self.start_segmentation)
         form_layout.addRow(seg_button)
 
+        full_button = QPushButton("Segment All")
+        full_button.clicked.connect(self.segment)
+        form_layout.addRow(full_button)
+
         self.trained = False
 
 
         self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=True)
 
 
+    def toggle_GPU(self):
+
+
+        self.use_gpu = self.GPU.isChecked()
 
 
     def toggle_foreground(self):
@@ -5869,7 +5934,6 @@ class MachineWindow(QMainWindow):
             self.fore_button.setChecked(True)
 
 
-
     def toggle_brush_mode(self):
         """Toggle brush mode on/off"""
         self.parent().brush_mode = self.brush_button.isChecked()
@@ -5889,34 +5953,164 @@ class MachineWindow(QMainWindow):
 
         self.brush_button.click()
 
-    def train_model(self):
+    def train_model(self, speed = True):
 
-        self.segmenter.train_batch(self.parent().channel_data[2], self.parent().channel_data[3])
+        self.kill_segmentation()
+        # Wait a bit for cleanup
+        time.sleep(0.1)
+        self.segmenter.train_batch(self.parent().channel_data[2], speed = speed, use_gpu = self.use_gpu)
         self.trained = True
 
+    def start_segmentation(self):
+
+        self.kill_segmentation()
+        print("Beginning new segmentation...")
+
+        if not self.trained:
+            return
+        else:
+            self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu)
+            self.segmentation_worker.chunk_processed.connect(self.update_display)  # Just update display
+            self.segmentation_worker.finished.connect(self.segmentation_finished)
+            current_xlim = self.parent().ax.get_xlim()
+            current_ylim = self.parent().ax.get_ylim()
+            self.segmenter.update_position(self.parent().current_slice, int((current_ylim[0] - current_ylim[1])/2), int((current_xlim[1] - current_xlim[0])/2))
+            self.segmentation_worker.start()
+
+
+
+    def update_display(self):
+        if not hasattr(self, '_last_update'):
+            self._last_update = 0
+        
+        current_time = time.time()
+        if current_time - self._last_update >= 1:  # Match worker's interval
+            try:
+                # Store current view state
+                current_xlim = self.parent().ax.get_xlim()
+                current_ylim = self.parent().ax.get_ylim()
+
+                self.segmenter.update_position(self.parent().current_slice, int((current_ylim[0] - current_ylim[1])/2), int((current_xlim[1] - current_xlim[0])/2))
+                if not self.parent().painting:
+                    # Only update if view limits are valid
+                    self.parent().update_display(preserve_zoom=(current_xlim, current_ylim))
+
+                    
+                    self._last_update = current_time
+            except Exception as e:
+                print(f"Display update error: {e}")
+
+    def segmentation_finished(self):
+        print("Segmentation completed")
+        current_xlim = self.parent().ax.get_xlim()
+        current_ylim = self.parent().ax.get_ylim()
+        self.parent().update_display(preserve_zoom=(current_xlim, current_ylim))
+        self.kill_segmentation()
+
+    def kill_segmentation(self):
+        if hasattr(self, 'segmentation_worker'):
+            self.segmentation_worker.stop()
+            del self.segmentation_worker  # Clean up reference
+
+
     def segment(self):
 
         if not self.trained:
             return
         else:
-            foreground_coords, background_coords = self.segmenter.segment_volume()
+            print("Segmenting entire volume with model...")
+            foreground_coords, background_coords = self.segmenter.segment_volume(gpu = self.use_gpu)
 
             # Clean up when done
             self.segmenter.cleanup()
 
         for z,y,x in foreground_coords:
-            self.parent().highlight_overlay[z,y,x] = True
+            self.parent().highlight_overlay[z,y,x] = 255
+
+        self.parent().load_channel(3, self.parent().highlight_overlay, True)
+
+        # Not exactly sure why we need all this but the channel buttons weren't loading like they normally do when load_channel() is called:
+        self.parent().channel_buttons[3].setEnabled(True)
+        self.parent().channel_buttons[3].click()
+        self.parent().delete_buttons[3].setEnabled(True)
+
+        self.parent().highlight_overlay = None
 
         self.parent().update_display()
 
+        print("Finished segmentation moved to Overlay 2. Use File -> Save(As) for disk saving.")
+
     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
+        self.kill_segmentation()
+
 
 
+class SegmentationWorker(QThread):
+    finished = pyqtSignal()
+    chunk_processed = pyqtSignal()
+    
+    def __init__(self, highlight_overlay, segmenter, use_gpu):
+        super().__init__()
+        self.overlay = highlight_overlay
+        self.segmenter = segmenter
+        self.use_gpu = use_gpu
+        self._stop = False
+        self.update_interval = 1  # Increased to 500ms
+        self.chunks_since_update = 0
+        self.chunks_per_update = 5  # Only update every 5 chunks
+        self.last_update = time.time()
+        
+    def stop(self):
+        self._stop = True
+        
+    def run(self):
+        try:
+            self.overlay.fill(False)
+            
+            for foreground_coords, background_coords in self.segmenter.segment_volume_realtime(gpu = self.use_gpu):
+                if self._stop:
+                    break
+                    
+                for z,y,x in foreground_coords:
+                    self.overlay[z,y,x] = 1
+                for z,y,x in background_coords:
+                    self.overlay[z,y,x] = 2
+                
+                # Update only after several chunks AND minimum time interval
+                self.chunks_since_update += 1
+                current_time = time.time()
+                if (self.chunks_since_update >= self.chunks_per_update and 
+                    current_time - self.last_update >= self.update_interval):
+                    self.chunk_processed.emit()
+                    self.chunks_since_update = 0
+                    self.last_update = current_time
+
+            self.finished.emit()
+            
+        except Exception as e:
+            print(f"Error in segmentation: {e}")
+            raise
+        
+    def run_batch(self):
+        try:
+            foreground_coords, _ = self.segmenter.segment_volume()
+            
+            # Modify the array directly
+            self.overlay.fill(False)
+            for z,y,x in foreground_coords:
+                self.overlay[z,y,x] = True
+                
+            self.finished.emit()
+            
+        except Exception as e:
+            print(f"Error in segmentation: {e}")
+            raise
+