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

nettracer3d/nettracer_gui.py

@@ -398,6 +398,7 @@ class ImageViewerWindow(QMainWindow):
         self.highlight_overlay = None
         self.highlight_bounds = None  # Store bounds for positioning
         self.mini_overlay = False # If the program is currently drawing the overlay by frame this will be true
+        self.mini_overlay_data = None #Actual data for mini overlay
         self.mini_thresh = (500*500*500) # Array volume to start using mini overlays for
 
     def start_left_scroll(self):
@@ -446,6 +447,7 @@ class ImageViewerWindow(QMainWindow):
         """
 
         self.mini_overlay = False #If this method is ever being called, it means we are rendering the entire overlay so mini overlay needs to reset.
+        self.mini_overlay_data = None
 
         def process_chunk(chunk_data, indices_to_check):
             """Process a single chunk of the array to create highlight mask"""
@@ -562,6 +564,8 @@ class ImageViewerWindow(QMainWindow):
 
     def create_highlight_overlay_slice(self, indices, bounds = False):
 
+        """Highlight overlay generation method specific for the segmenter interactive mode"""
+
 
         def process_chunk_bounds(chunk_data, indices_to_check):
             """Process a single chunk of the array to create highlight mask"""
@@ -638,7 +642,7 @@ class ImageViewerWindow(QMainWindow):
     def create_mini_overlay(self, node_indices = None, edge_indices = None):
 
         """
-        Create a binary overlay highlighting specific nodes and/or edges using parallel processing, one slice at a time for efficiency.
+        Create a highlight overlay one slice at a time.
         
         Args:
             node_indices (list): List of node indices to highlight
@@ -668,28 +672,29 @@ class ImageViewerWindow(QMainWindow):
         current_xlim = self.ax.get_xlim() if hasattr(self, 'ax') and self.ax.get_xlim() != (0, 1) else None
         current_ylim = self.ax.get_ylim() if hasattr(self, 'ax') and self.ax.get_ylim() != (0, 1) else None
         
-        if not node_indices and not edge_indices:
-            self.highlight_overlay = None
+        if not node_indices and not edge_indices: #Theoretically this can't be called because it uses full highlight overlay method for empty clicks
+            self.mini_overlay_data = None
+            self.mini_overlay = False
             self.update_display(preserve_zoom=(current_xlim, current_ylim))
             return
             
-        # Get the shape of the full array from any existing channel
+        # Get the shape of the mini array from any existing channel
         for channel in self.channel_data:
             if channel is not None:
                 full_shape = channel.shape
+                full_shape = (full_shape[1], full_shape[2]) #Just get (Y, X) shape
                 break
         else:
             return  # No valid channels to get shape from
             
         # Initialize full-size overlay
-        if self.highlight_overlay is None:
-            self.highlight_overlay = np.zeros(full_shape, dtype=np.uint8)
+        self.mini_overlay_data = np.zeros(full_shape, dtype=np.uint8)
         
         # Get number of CPU cores
         num_cores = mp.cpu_count()
         
         # Calculate chunk size along y-axis
-        chunk_size = full_shape[1] // num_cores
+        chunk_size = full_shape[0] // num_cores
         if chunk_size < 1:
             chunk_size = 1
         
@@ -699,8 +704,8 @@ class ImageViewerWindow(QMainWindow):
                 
             # Create chunks
             chunks = []
-            for i in range(0, array_shape[1], chunk_size):
-                end = min(i + chunk_size, array_shape[1])
+            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
@@ -729,11 +734,10 @@ class ImageViewerWindow(QMainWindow):
                 edge_overlay = None
                         
         # Combine results
-        self.highlight_overlay[self.current_slice, :, :] = np.zeros_like(self.highlight_overlay[self.current_slice, :, :])
         if node_overlay is not None:
-            self.highlight_overlay[self.current_slice, :, :] = np.maximum(self.highlight_overlay[self.current_slice, :, :], node_overlay)
+            self.mini_overlay_data = np.maximum(self.mini_overlay_data, node_overlay)
         if edge_overlay is not None:
-            self.highlight_overlay[self.current_slice, :, :] = np.maximum(self.highlight_overlay[self.current_slice, :, :], edge_overlay)
+            self.mini_overlay_data = np.maximum(self.mini_overlay_data, edge_overlay)
 
                 
         # Update display
@@ -3488,6 +3492,8 @@ class ImageViewerWindow(QMainWindow):
             if self.mini_overlay == True: #If we are rendering the highlight overlay for selected values one at a time.
                 self.create_mini_overlay(node_indices = self.clicked_values['nodes'], edge_indices = self.clicked_values['edges'])
             self.update_display(preserve_zoom=view_settings)
+            if self.machine_window is not None:
+                self.machine_window.poke_segmenter()
             self.pending_slice = None
 
     def update_brightness(self, channel_index, values):
@@ -3623,7 +3629,15 @@ class ImageViewerWindow(QMainWindow):
             self.create_highlight_overlay_slice(self.targs, bounds = self.bounds)
 
         # Add highlight overlay if it exists
-        if self.highlight_overlay is not None and self.highlight and self.machine_window is None:
+        if self.mini_overlay and self.highlight and self.machine_window is None:
+            highlight_cmap = LinearSegmentedColormap.from_list(
+                'highlight',
+                [(0, 0, 0, 0), (1, 1, 0, 1)]  # yellow
+            )
+            self.ax.imshow(self.mini_overlay_data,
+                         cmap=highlight_cmap,
+                         alpha=0.5)
+        elif 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',
@@ -6202,7 +6216,18 @@ class MachineWindow(QMainWindow):
         self.GPU.setChecked(False)
         self.GPU.clicked.connect(self.toggle_GPU)
         self.use_gpu = False
+        self.two = QPushButton("Train By 2D Slice Patterns (Cheaper - CPU only)")
+        self.two.setCheckable(True)
+        self.two.setChecked(True)
+        self.two.clicked.connect(self.toggle_two)
+        self.use_two = True
+        self.three = QPushButton("Train by 3D Patterns")
+        self.three.setCheckable(True)
+        self.three.setChecked(False)
+        self.three.clicked.connect(self.toggle_three)
         processing_layout.addWidget(self.GPU)
+        processing_layout.addWidget(self.two)
+        processing_layout.addWidget(self.three)
         processing_group.setLayout(processing_layout)
 
         # Group 3: Training Options
@@ -6220,6 +6245,7 @@ class MachineWindow(QMainWindow):
         segmentation_group = QGroupBox("Segmentation")
         segmentation_layout = QVBoxLayout()
         seg_button = QPushButton("Preview Segment")
+        self.seg_button = seg_button
         seg_button.clicked.connect(self.start_segmentation)
         full_button = QPushButton("Segment All")
         full_button.clicked.connect(self.segment)
@@ -6237,9 +6263,12 @@ class MachineWindow(QMainWindow):
         self.setCentralWidget(main_widget)
 
         self.trained = False
+        self.previewing = False
 
 
         self.segmenter = segmenter.InteractiveSegmenter(active_data, use_gpu=True)
+        self.segmentation_worker = None
+
 
 
     def toggle_GPU(self):
@@ -6247,6 +6276,27 @@ class MachineWindow(QMainWindow):
 
         self.use_gpu = self.GPU.isChecked()
 
+    def toggle_two(self):
+        if self.two.isChecked():
+            # If button two is checked, ensure button three is unchecked
+            self.three.setChecked(False)
+            self.use_two = True
+        else:
+            # If button three is checked, ensure button two is unchecked
+            self.three.setChecked(True)
+            self.use_two = False
+
+    def toggle_three(self):
+        if self.three.isChecked():
+            # If button two is checked, ensure button three is unchecked
+            self.two.setChecked(False)
+            self.use_two = False
+        else:
+            # If button three is checked, ensure button two is unchecked
+            self.two.setChecked(True)
+            self.use_two = True
+
+
 
     def toggle_foreground(self):
 
@@ -6295,15 +6345,29 @@ class MachineWindow(QMainWindow):
         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
+        if not self.use_two:
+            self.previewing = False
+        try:
+            self.segmenter.train_batch(self.parent().channel_data[2], speed = speed, use_gpu = self.use_gpu, use_two = self.use_two)
+            self.trained = True
+        except MemoryError:
+            QMessageBox.critical(
+                self,
+                "Alert",
+                "Out of memory computing feature maps. Note these for 3D require 7x the RAM of the active image (or 9x for the detailed map).\n Please use 2D slice models if you do not have enough RAM."
+            )
+
 
     def start_segmentation(self):
 
         self.kill_segmentation()
         time.sleep(0.1)
 
-        print("Beginning new segmentation...")
+        if self.use_two:
+            self.previewing = True
+        else:
+            print("Beginning new segmentation...")
+
 
         if self.parent().active_channel == 0:
             if self.parent().channel_data[0] is not None:
@@ -6317,7 +6381,7 @@ class MachineWindow(QMainWindow):
         if not self.trained:
             return
         else:
-            self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu)
+            self.segmentation_worker = SegmentationWorker(self.parent().highlight_overlay, self.segmenter, self.use_gpu, self.use_two, self.previewing, self)
             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()
@@ -6345,10 +6409,42 @@ class MachineWindow(QMainWindow):
         msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
         return msg.exec() == QMessageBox.StandardButton.Yes
 
+
+
+    def check_for_z_change(self):
+        current_z = self.parent().current_slice
+        if not hasattr(self, '_last_z'):
+            self._last_z = current_z
+            return False
+        
+        changed = (self._last_z != current_z)
+        self._last_z = current_z
+        
+        if changed and self.previewing and self.segmentation_worker is not None:
+            self.segmentation_worker.stop()
+            time.sleep(0.1)
+            
+            # Force regeneration of chunks
+            self.segmenter.realtimechunks = None
+            
+            # Restart the worker
+            self.start_segmentation()
+            
+        return changed
+
     def update_display(self):
         if not hasattr(self, '_last_update'):
             self._last_update = 0
         
+        current_z = self.parent().current_slice
+        if not hasattr(self, '_last_z'):
+            self._last_z = current_z
+
+        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
             try:
@@ -6366,19 +6462,84 @@ class MachineWindow(QMainWindow):
             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()
+                
+            except Exception as e:
+                print(f"Error in poke_segmenter: {e}")
+                import traceback
+                traceback.print_exc()
+
     def segmentation_finished(self):
-        print("Segmentation completed")
+        if not self.use_two:
+            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))
+        
+        # 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 = {}
+                
+            current_time = time.time()
+            
+            # 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
+                    
+                # Create a new worker after a brief delay
+                QTimer.singleShot(500, self.start_segmentation)
+
 
 
     def kill_segmentation(self):
-        if hasattr(self, 'segmentation_worker'):
+        if hasattr(self, 'segmentation_worker') and self.segmentation_worker is not None:
+            # Signal the thread to stop
             self.segmentation_worker.stop()
-            del self.segmentation_worker  # Clean up reference
+            
+            # Wait for the thread to finish
+            if self.segmentation_worker.isRunning():
+                self.segmentation_worker.wait(1000)  # Wait up to 1 second
+                
+                # If thread is still running after timeout, try to force termination
+                if self.segmentation_worker.isRunning():
+                    self.segmentation_worker.terminate()
+                    self.segmentation_worker.wait()  # Wait for it to be terminated
+            
+            # Now safe to delete
+            del self.segmentation_worker
+            self.segmentation_worker = None
 
 
     def segment(self):
@@ -6391,6 +6552,8 @@ class MachineWindow(QMainWindow):
             self.kill_segmentation()
             time.sleep(0.1)
 
+            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]
@@ -6420,20 +6583,26 @@ class MachineWindow(QMainWindow):
 
         self.parent().update_display()
 
+        self.previewing = False
+
         print("Finished segmentation moved to Overlay 2. Use File -> Save(As) for disk saving.")
 
     def closeEvent(self, event):
         if self.parent().isVisible():
             if self.confirm_close_dialog():
-
+                # Clean up resources before closing
                 if self.brush_button.isChecked():
                     self.silence_button()
                     self.toggle_brush_mode()
+                
                 self.parent().pen_button.setEnabled(True)
                 self.parent().brush_mode = False
-                self.parent().machine_window = None
+                
+                # Kill the segmentation thread and wait for it to finish
                 self.kill_segmentation()
-                time.sleep(0.1)
+                time.sleep(0.2)  # Give additional time for cleanup
+                
+                self.parent().machine_window = None
             else:
                 event.ignore()
 
@@ -6444,48 +6613,84 @@ class SegmentationWorker(QThread):
     finished = pyqtSignal()
     chunk_processed = pyqtSignal()
     
-    def __init__(self, highlight_overlay, segmenter, use_gpu):
+    def __init__(self, highlight_overlay, segmenter, use_gpu, use_two, previewing, machine_window):
         super().__init__()
         self.overlay = highlight_overlay
         self.segmenter = segmenter
         self.use_gpu = use_gpu
+        self.use_two = use_two
+        self.previewing = previewing
+        self.machine_window = machine_window
         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.poked = False # If it should wake up or not
         self.last_update = time.time()
         
     def stop(self):
         self._stop = True
+
+    def get_poked(self):
+        self.poked = 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
+            # 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):
+                    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
-
+                    # 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
+                
+            else:
+                # Original 3D approach
+                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
+                    
+                    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
-        
+            import traceback
+            traceback.print_exc()
+
     def run_batch(self):
         try:
             foreground_coords, _ = self.segmenter.segment_volume()