coralnet-toolbox 0.0.75__py2.py3-none-any.whl → 0.0.76__py2.py3-none-any.whl

coralnet_toolbox/QtSystemMonitor.py

@@ -0,0 +1,370 @@
+import psutil
+import GPUtil
+import collections
+import time
+import platform
+
+import pyqtgraph as pg
+from PyQt5.QtWidgets import QMainWindow, QGridLayout, QWidget, QLabel
+from PyQt5.QtCore import Qt, QTimer
+from PyQt5.QtGui import QFont
+
+from coralnet_toolbox.Icons import get_icon
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Classes
+# ----------------------------------------------------------------------------------------------------------------------
+
+
+class SystemMonitor(QMainWindow):
+    """
+    A PyQt5 application to monitor and display real-time CPU, Memory, and GPU usage.
+    The monitoring starts/stops automatically when the window is shown/hidden.
+    """
+    def __init__(self):
+        super().__init__()
+
+        # --- Window Properties ---
+        self.setWindowTitle("System Monitor")
+        self.setWindowIcon(get_icon("system_monitor.png"))
+        self.setGeometry(100, 100, 800, 900)
+
+        # --- Data Storage ---
+        # Use collections.deque for efficient, fixed-size data storage
+        self.history_size = 100  # Number of data points to display
+        self.cpu_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.mem_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.gpu_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        
+        # Per-core CPU data
+        self.cpu_cores = psutil.cpu_count(logical=True)
+        self.per_core_data = []
+        for _ in range(self.cpu_cores):
+            self.per_core_data.append(collections.deque([0] * self.history_size, maxlen=self.history_size))
+
+        # Disk I/O data
+        self.disk_read_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.disk_write_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.prev_disk_io = psutil.disk_io_counters()
+        self.prev_time = time.monotonic()
+        
+        # Network traffic data
+        self.net_sent_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.net_recv_data = collections.deque([0] * self.history_size, maxlen=self.history_size)
+        self.prev_net_io = psutil.net_io_counters()
+        
+        # X-axis data (time steps)
+        self.x_axis = list(range(self.history_size))
+
+        # --- Initialize UI ---
+        self.setup_layout()
+
+        # --- Timer for Real-Time Updates ---
+        # Set up a QTimer to trigger the update_plots method every 1000 ms (1 second)
+        # We do NOT start the timer here. It will be started by the showEvent.
+        self.timer = QTimer()
+        self.timer.setInterval(1000)
+        self.timer.timeout.connect(self.update_plots)
+
+    def setup_layout(self):
+        """
+        Initializes the user interface, setting up the layout and plots.
+        """
+        # --- Central Widget and Layout ---
+        central_widget = QWidget()
+        self.setCentralWidget(central_widget)
+        layout = QGridLayout(central_widget)
+        layout.setVerticalSpacing(20)  # Add more vertical spacing between plots
+
+        # --- Styling ---
+        pg.setConfigOption('background', 'w')  # White background for graphs
+        pg.setConfigOption('foreground', '#000000')  # Black foreground elements in graphs
+        plot_label_style = {"color": "#000000", "font-size": "14pt"}  # Black text for plot labels
+        title_style = {"color": "#000000", "font-size": "16pt", "font-weight": "bold"}  # Black text
+
+        # --- Stats Labels - Only Hardware Info ---
+        self.stats_widget = QWidget()
+        stats_layout = QGridLayout(self.stats_widget)
+        stats_layout.setHorizontalSpacing(20)  # Add horizontal spacing between items
+        stats_layout.setContentsMargins(10, 10, 10, 20)  # Add margins around the widget
+
+        # Get system hardware information
+        cpu_info = platform.processor() or "CPU"
+        total_memory = round(psutil.virtual_memory().total / (1024**3), 1)
+        
+        try:
+            gpus = GPUtil.getGPUs()
+            gpu_info = gpus[0].name if gpus else "No GPU detected"
+        except Exception:
+            gpu_info = "GPU information unavailable"
+        
+        # Hardware info labels with consistent styling and alignment
+        info_label_style = "color: black; font-size: 11pt; font-weight: bold; padding: 5px;"
+
+        self.cpu_info_label = QLabel(f"CPU: {cpu_info}")
+        self.cpu_info_label.setStyleSheet(info_label_style)
+        self.cpu_info_label.setAlignment(Qt.AlignCenter)
+        stats_layout.addWidget(self.cpu_info_label, 0, 0)
+        
+        self.mem_info_label = QLabel(f"Memory: {total_memory:.1f} GB")
+        self.mem_info_label.setStyleSheet(info_label_style)
+        self.mem_info_label.setAlignment(Qt.AlignCenter)
+        stats_layout.addWidget(self.mem_info_label, 0, 1)
+        
+        self.gpu_info_label = QLabel(f"GPU: {gpu_info}")
+        self.gpu_info_label.setStyleSheet(info_label_style)
+        self.gpu_info_label.setAlignment(Qt.AlignCenter)
+        stats_layout.addWidget(self.gpu_info_label, 0, 2)
+        
+        # Set column stretching to distribute space evenly
+        stats_layout.setColumnStretch(0, 1)
+        stats_layout.setColumnStretch(1, 1)
+        stats_layout.setColumnStretch(2, 1)
+
+        layout.addWidget(self.stats_widget, 0, 0, 1, 2)
+
+        # Helper function to add bottom space for x-axis and enable auto-range
+        def configure_plot(plot_widget, fixed_y_range=None):
+            plot_widget.getPlotItem().getAxis('bottom').setHeight(40)
+            plot_widget.getPlotItem().layout.setContentsMargins(10, 10, 10, 30)
+            
+            # If fixed_y_range is None, enable auto-ranging, otherwise set fixed range
+            if fixed_y_range is None:
+                plot_widget.enableAutoRange(axis='y')
+            else:
+                min_val, max_val = fixed_y_range
+                plot_widget.setYRange(min_val, max_val, padding=0.05)
+                
+            return plot_widget
+            
+        # Helper function to create a value label for a plot
+        def create_value_label(plot_widget, initial_text="Current: N/A"):
+            text_item = pg.TextItem(text=initial_text, color='#000000', anchor=(0, 0))
+            font = QFont()
+            font.setBold(True)
+            text_item.setFont(font)
+            plot_widget.addItem(text_item)
+            # Position at top-left corner with some margin
+            text_item.setPos(10, 10)
+            return text_item
+
+        # --- CPU Plot --- (fixed 0-100 range)
+        self.cpu_plot_widget = pg.PlotWidget()
+        self.cpu_plot_widget.setTitle("CPU Usage (%)", **title_style)
+        self.cpu_plot_widget.setLabel("left", "Usage", units="%", **plot_label_style)
+        self.cpu_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.cpu_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        self.cpu_curve = self.cpu_plot_widget.plot(self.x_axis, list(self.cpu_data), pen=pg.mkPen('#00A3FF', width=3))
+        self.cpu_value_label = create_value_label(self.cpu_plot_widget, "Current: 0%")
+        configure_plot(self.cpu_plot_widget, fixed_y_range=(0, 100))
+        layout.addWidget(self.cpu_plot_widget, 1, 0)
+
+        # --- Per-Core CPU Plot --- (fixed 0-100 range)
+        self.core_plot_widget = pg.PlotWidget()
+        self.core_plot_widget.setTitle("Per-Core CPU Usage (%)", **title_style)
+        self.core_plot_widget.setLabel("left", "Usage", units="%", **plot_label_style)
+        self.core_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.core_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        
+        self.per_core_curves = []
+        for i in range(self.cpu_cores):
+            color = pg.intColor(i, hues=self.cpu_cores)
+            self.per_core_curves.append(
+                self.core_plot_widget.plot(
+                    self.x_axis, 
+                    list(self.per_core_data[i]), 
+                    pen=pg.mkPen(color, width=2)
+                )
+            )
+        configure_plot(self.core_plot_widget, fixed_y_range=(0, 100))
+        layout.addWidget(self.core_plot_widget, 1, 1)
+
+        # --- Memory Plot --- (fixed 0-100 range)
+        self.mem_plot_widget = pg.PlotWidget()
+        self.mem_plot_widget.setTitle("Memory Usage (%)", **title_style)
+        self.mem_plot_widget.setLabel("left", "Usage", units="%", **plot_label_style)
+        self.mem_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.mem_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        self.mem_curve = self.mem_plot_widget.plot(self.x_axis, list(self.mem_data), pen=pg.mkPen('#FF5733', width=3))
+        self.mem_value_label = create_value_label(self.mem_plot_widget, "Current: 0%")
+        configure_plot(self.mem_plot_widget, fixed_y_range=(0, 100))
+        layout.addWidget(self.mem_plot_widget, 2, 0)
+
+        # --- Disk I/O Plot --- (auto-ranging)
+        self.disk_plot_widget = pg.PlotWidget()
+        self.disk_plot_widget.setTitle("Disk I/O (MB/s)", **title_style)
+        self.disk_plot_widget.setLabel("left", "Transfer Rate", units="MB/s", **plot_label_style)
+        self.disk_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.disk_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        
+        self.disk_read_curve = self.disk_plot_widget.plot(
+            self.x_axis, 
+            list(self.disk_read_data), 
+            pen=pg.mkPen('#FFDD33', width=3)
+        )
+        self.disk_write_curve = self.disk_plot_widget.plot(
+            self.x_axis, 
+            list(self.disk_write_data), 
+            pen=pg.mkPen('#33DDFF', width=3)
+        )
+        
+        self.disk_value_label = create_value_label(self.disk_plot_widget, "R: 0 MB/s, W: 0 MB/s")
+        configure_plot(self.disk_plot_widget, fixed_y_range=None)  # Auto-range
+        layout.addWidget(self.disk_plot_widget, 2, 1)
+
+        # --- GPU Plot --- (fixed 0-100 range)
+        self.gpu_plot_widget = pg.PlotWidget()
+        self.gpu_plot_widget.setTitle("GPU Usage (%)", **title_style)
+        self.gpu_plot_widget.setLabel("left", "Usage", units="%", **plot_label_style)
+        self.gpu_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.gpu_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        self.gpu_curve = self.gpu_plot_widget.plot(self.x_axis, list(self.gpu_data), pen=pg.mkPen('#33FF57', width=3))
+        self.gpu_value_label = create_value_label(self.gpu_plot_widget, "Current: 0%")
+        configure_plot(self.gpu_plot_widget, fixed_y_range=(0, 100))
+        layout.addWidget(self.gpu_plot_widget, 3, 0)
+
+        # --- Network Traffic Plot --- (auto-ranging)
+        self.net_plot_widget = pg.PlotWidget()
+        self.net_plot_widget.setTitle("Network Traffic (MB/s)", **title_style)
+        self.net_plot_widget.setLabel("left", "Transfer Rate", units="MB/s", **plot_label_style)
+        self.net_plot_widget.setLabel("bottom", "Time (s)", **plot_label_style)
+        self.net_plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        
+        self.net_sent_curve = self.net_plot_widget.plot(
+            self.x_axis, 
+            list(self.net_sent_data), 
+            pen=pg.mkPen('#FF33DD', width=3)
+        )
+        self.net_recv_curve = self.net_plot_widget.plot(
+            self.x_axis, 
+            list(self.net_recv_data), 
+            pen=pg.mkPen('#DD33FF', width=3)
+        )
+        
+        self.net_value_label = create_value_label(self.net_plot_widget, "Up: 0 MB/s, Down: 0 MB/s")
+        configure_plot(self.net_plot_widget, fixed_y_range=None)  # Auto-range
+        layout.addWidget(self.net_plot_widget, 3, 1)
+
+    def update_plots(self):
+        """
+        Fetches new system stats and updates the plot data.
+        This method is called by the QTimer.
+        """
+        curr_time = time.monotonic()
+        time_delta = curr_time - self.prev_time
+        
+        # --- Get CPU, Memory and GPU percentages ---
+        # CPU usage
+        cpu_percent = psutil.cpu_percent()
+        self.cpu_data.append(cpu_percent)
+        self.cpu_value_label.setText(f"Current: {cpu_percent:.1f}%")
+        
+        # Per-core CPU usage
+        per_core = psutil.cpu_percent(percpu=True)
+        for i, usage in enumerate(per_core):
+            if i < len(self.per_core_data):
+                self.per_core_data[i].append(usage)
+                self.per_core_curves[i].setData(self.x_axis, list(self.per_core_data[i]))
+
+        # Memory usage
+        mem_percent = psutil.virtual_memory().percent
+        self.mem_data.append(mem_percent)
+        self.mem_value_label.setText(f"Current: {mem_percent:.1f}%")
+
+        # GPU usage
+        try:
+            gpus = GPUtil.getGPUs()
+            if gpus:
+                gpu_percent = gpus[0].load * 100
+                self.gpu_value_label.setText(f"Current: {gpu_percent:.1f}%")
+            else:
+                gpu_percent = 0
+                self.gpu_value_label.setText("Current: 0%")
+        except Exception:
+            gpu_percent = 0
+            self.gpu_value_label.setText("Current: 0%")
+
+        self.gpu_data.append(gpu_percent)
+        
+        # Disk I/O
+        curr_disk_io = psutil.disk_io_counters()
+        
+        # Calculate read/write rates in MB/s
+        if curr_disk_io and self.prev_disk_io:
+            read_bytes = curr_disk_io.read_bytes - self.prev_disk_io.read_bytes
+            write_bytes = curr_disk_io.write_bytes - self.prev_disk_io.write_bytes
+            read_mb_per_sec = (read_bytes / time_delta) / (1024**2)
+            write_mb_per_sec = (write_bytes / time_delta) / (1024**2)
+            
+            self.disk_read_data.append(read_mb_per_sec)
+            self.disk_write_data.append(write_mb_per_sec)
+            
+            # Update the value label with current values
+            self.disk_value_label.setText(f"R: {read_mb_per_sec:.1f} MB/s, W: {write_mb_per_sec:.1f} MB/s")
+            
+            # Let pyqtgraph handle auto-ranging automatically
+            self.disk_plot_widget.enableAutoRange(axis='y')
+    
+        self.prev_disk_io = curr_disk_io
+        
+        # Network traffic
+        curr_net_io = psutil.net_io_counters()
+        if curr_net_io and self.prev_net_io:
+            sent_bytes = curr_net_io.bytes_sent - self.prev_net_io.bytes_sent
+            recv_bytes = curr_net_io.bytes_recv - self.prev_net_io.bytes_recv
+            
+            sent_mb_per_sec = (sent_bytes / time_delta) / (1024**2)
+            recv_mb_per_sec = (recv_bytes / time_delta) / (1024**2)
+            
+            self.net_sent_data.append(sent_mb_per_sec)
+            self.net_recv_data.append(recv_mb_per_sec)
+            
+            # Update the value label with current values
+            self.net_value_label.setText(f"Up: {sent_mb_per_sec:.1f} MB/s, Down: {recv_mb_per_sec:.1f} MB/s")
+            
+            # Let pyqtgraph handle auto-ranging automatically
+            self.net_plot_widget.enableAutoRange(axis='y')
+    
+        self.prev_net_io = curr_net_io
+
+        # --- Update Plot Curves ---
+        # Convert deque to list for plotting
+        self.cpu_curve.setData(self.x_axis, list(self.cpu_data))
+        self.mem_curve.setData(self.x_axis, list(self.mem_data))
+        self.gpu_curve.setData(self.x_axis, list(self.gpu_data))
+        self.disk_read_curve.setData(self.x_axis, list(self.disk_read_data))
+        self.disk_write_curve.setData(self.x_axis, list(self.disk_write_data))
+        self.net_sent_curve.setData(self.x_axis, list(self.net_sent_data))
+        self.net_recv_curve.setData(self.x_axis, list(self.net_recv_data))
+
+        self.prev_time = curr_time
+
+    # --- Event Handlers to Control Monitoring ---
+    
+    def showEvent(self, event):
+        """
+        Overrides the QWidget's showEvent.
+        Starts the timer when the window is shown.
+        """
+        super().showEvent(event)
+        if not self.timer.isActive():
+            self.timer.start()
+
+    def hideEvent(self, event):
+        """
+        Overrides the QWidget's hideEvent.
+        Stops the timer when the window is hidden (e.g., minimized).
+        """
+        super().hideEvent(event)
+        if self.timer.isActive():
+            self.timer.stop()
+
+    def closeEvent(self, event):
+        """
+        Overrides the QMainWindow's closeEvent.
+        Ensures the timer is stopped when the window is closed.
+        """
+        if self.timer.isActive():
+            self.timer.stop()
+        super().closeEvent(event)

coralnet_toolbox/Results/ConvertResults.py

@@ -94,7 +94,7 @@ class ConvertResults:
             names: Dictionary mapping class ids to class names (optional)
 
         Returns:
-            generator: Yields Ultralytics Results objects
+            list: List of Ultralytics Results objects
         """
         # Convert single inputs to lists
         if not isinstance(detections, list):
@@ -108,8 +108,15 @@ class ConvertResults:
         if not image_paths:
             image_paths = [None] * len(images)
 
+        results_list = []
         for detection, image, path in zip(detections, images, image_paths):
             # Ensure image is numpy array
+            if isinstance(image, str):
+                # If image is a string (path), we can't proceed because we need the actual image data
+                raise ValueError(
+                    f"Expected image array, got path string: {image}. Please load the image before passing."
+                )
+                
             if torch.is_tensor(image):
                 image = image.cpu().numpy()
 
@@ -118,7 +125,8 @@ class ConvertResults:
                 names = {i: str(i) for i in range(len(detection))} if len(detection) > 0 else {}
 
             if len(detection) == 0:
-                return [Results(orig_img=image, path=path, names=names)]
+                results_list.append(Results(orig_img=image, path=path, names=names))
+                continue
 
             # Handle masks if present
             if hasattr(detection, 'mask') and detection.mask is not None:
@@ -142,11 +150,9 @@ class ConvertResults:
                 scaled_boxes = scaled_boxes.unsqueeze(0)
             scaled_boxes = torch.cat([scaled_boxes, scores, cls], dim=1)
 
-            # Create and return Results object
-            return [Results(image,
-                            path=path,
-                            names=names,
-                            boxes=scaled_boxes, 
-                            masks=scaled_masks)]
+            # Create Results object
+            results_list.append(Results(image, path=path, names=names, boxes=scaled_boxes, masks=scaled_masks))
+
+        return results_list
     
 

coralnet_toolbox/Results/ResultsProcessor.py

@@ -1,6 +1,7 @@
 from PyQt5.QtCore import QPointF
 
 from torchvision.ops import nms
+from ultralytics.utils.nms import TorchNMS
 
 from coralnet_toolbox.Annotations.QtPolygonAnnotation import PolygonAnnotation
 from coralnet_toolbox.Annotations.QtRectangleAnnotation import RectangleAnnotation
@@ -62,7 +63,7 @@ class ResultsProcessor:
     def filter_by_iou(self, results):
         """Filter the results based on the IoU threshold."""
         try:
-            results = results[nms(results.boxes.xyxy, results.boxes.conf, self.iou_thresh)]
+            results = results[TorchNMS.fast_nms(results.boxes.xyxy, results.boxes.conf, self.iou_thresh)]
         except Exception as e:
             print(f"Warning: Failed to filter results by IoU\n{e}")
 
@@ -106,7 +107,7 @@ class ResultsProcessor:
         Get the indices of results that pass the IoU threshold.
         """
         try:
-            indices = nms(results.boxes.xyxy, results.boxes.conf, self.iou_thresh).tolist()
+            indices = TorchNMS.fast_nms(results.boxes.xyxy, results.boxes.conf, self.iou_thresh).tolist()
         except Exception as e:
             print(f"Warning: Failed to get indices for IoU\n{e}")
             indices = []

coralnet_toolbox/SAM/QtDeployGenerator.py

@@ -36,7 +36,7 @@ class DeployGeneratorDialog(QDialog):
 
     def __init__(self, main_window, parent=None):
         """
-        Initialize the AutoDistillDeployModelDialog.
+        Initialize the DeployGeneratorDialog.
 
         Args:
             main_window: The main application window.

coralnet_toolbox/SAM/QtDeployPredictor.py

@@ -143,6 +143,12 @@ class DeployPredictorDialog(QDialog):
         group_box = QGroupBox("Parameters")
         layout = QFormLayout()
         
+        # Output type dropdown (polygon or rectangle)
+        self.output_type_dropdown = QComboBox()
+        self.output_type_dropdown.addItems(["Polygon", "Rectangle"])
+        self.output_type_dropdown.setCurrentIndex(0)  # Default to Polygon
+        layout.addRow("Output Type:", self.output_type_dropdown)
+        
         # Allow holes dropdown
         self.allow_holes_dropdown = QComboBox()
         self.allow_holes_dropdown.addItems(["True", "False"])
@@ -246,6 +252,10 @@ class DeployPredictorDialog(QDialog):
         """Return the current setting for allowing holes."""
         return self.allow_holes_dropdown.currentText() == "True"
     
+    def get_output_type(self):
+        """Return the current setting for output type."""
+        return self.output_type_dropdown.currentText()
+    
     def initialize_uncertainty_threshold(self):
         """Initialize the uncertainty threshold slider with the current value"""
         current_value = self.main_window.get_uncertainty_thresh()

coralnet_toolbox/SeeAnything/QtDeployGenerator.py

@@ -10,10 +10,10 @@ import torch
 from torch.cuda import empty_cache
 
 import pyqtgraph as pg
-from pyqtgraph.Qt import QtGui
 
 from ultralytics import YOLOE
 from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
+from ultralytics.models.yolo.yoloe import YOLOEVPDetectPredictor
 
 from PyQt5.QtCore import Qt
 from PyQt5.QtWidgets import (QMessageBox, QVBoxLayout, QApplication, QFileDialog,
@@ -90,7 +90,7 @@ class DeployGeneratorDialog(QDialog):
         self.imported_vpes = []  # VPEs loaded from file
         self.reference_vpes = []  # VPEs created from reference images
 
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = None  # Will be set in showEvent
 
         # Main vertical layout for the dialog
         self.layout = QVBoxLayout(self)
@@ -195,6 +195,8 @@ class DeployGeneratorDialog(QDialog):
         self.initialize_iou_threshold()
         self.initialize_area_threshold()
         
+        # Update the device
+        self.device = self.main_window.device
         # Configure the image window's UI elements for this specific dialog
         self.configure_image_window_for_dialog()
         # Sync with main window's images BEFORE updating labels
@@ -767,20 +769,20 @@ class DeployGeneratorDialog(QDialog):
         try:
             # Load the VPE file
             loaded_data = torch.load(file_path)
-            
-            # TODO Move tensors to the appropriate device
-            # device = self.main_window.device
+
+            # Move tensors to the appropriate device
+            device = self.main_window.device
             
             # Check format type and handle appropriately
             if isinstance(loaded_data, list):
                 # New format: list of VPE tensors
-                self.imported_vpes = [vpe.to(self.device) for vpe in loaded_data]
+                self.imported_vpes = [vpe.to(device) for vpe in loaded_data]
                 vpe_count = len(self.imported_vpes)
                 self.status_bar.setText(f"Loaded {vpe_count} VPE tensors from file")
                 
             elif isinstance(loaded_data, torch.Tensor):
                 # Legacy format: single tensor - convert to list for consistency
-                loaded_vpe = loaded_data.to(self.device)
+                loaded_vpe = loaded_data.to(device)
                 # Store as a single-item list
                 self.imported_vpes = [loaded_vpe]
                 self.status_bar.setText("Loaded 1 VPE tensor from file (legacy format)")
@@ -948,7 +950,7 @@ class DeployGeneratorDialog(QDialog):
         self.model_path = self.model_combo.currentText()
 
         # Load model using registry
-        self.loaded_model = YOLOE(self.model_path, verbose=False).to(self.device)  # TODO
+        self.loaded_model = YOLOE(self.model_path, verbose=False).to(self.device)
 
         # Create a dummy visual dictionary for standard model loading
         visual_prompts = dict(
@@ -966,7 +968,7 @@ class DeployGeneratorDialog(QDialog):
         self.loaded_model.predict(
             np.zeros((640, 640, 3), dtype=np.uint8),
             visual_prompts=visual_prompts.copy(),  # This needs to happen to properly initialize the predictor
-            predictor=YOLOEVPSegPredictor,  # This also needs to be SegPredictor, no matter what
+            predictor=YOLOEVPDetectPredictor if self.task == "detect" else YOLOEVPSegPredictor,
             imgsz=640,
             conf=0.99,
         )
@@ -1108,6 +1110,9 @@ class DeployGeneratorDialog(QDialog):
         # We will predict on the first work area/full image.
         input_image = inputs[0] 
 
+        # Set the predictor
+        predictor = YOLOEVPDetectPredictor if self.task == "detect" else YOLOEVPSegPredictor
+
         # Iterate through each reference image and its annotations
         for ref_path, ref_annotations in reference_dict.items():
             # The 'refer_image' parameter is the path to the current reference image
@@ -1123,7 +1128,7 @@ class DeployGeneratorDialog(QDialog):
             results = self.loaded_model.predict(input_image,
                                                 refer_image=ref_path,
                                                 visual_prompts=visual_prompts,
-                                                predictor=YOLOEVPSegPredictor,  # TODO This is necessary here?
+                                                predictor=predictor,
                                                 imgsz=self.imgsz_spinbox.value(),
                                                 conf=self.main_window.get_uncertainty_thresh(),
                                                 iou=self.main_window.get_iou_thresh(),

coralnet_toolbox/SeeAnything/QtDeployPredictor.py

@@ -1,6 +1,5 @@
 import warnings
 
-import os
 import gc
 
 import numpy as np
@@ -14,7 +13,6 @@ from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
 from ultralytics.models.yolo.yoloe import YOLOEVPDetectPredictor
 
 from PyQt5.QtCore import Qt
-from PyQt5.QtGui import QColor
 from PyQt5.QtWidgets import (QApplication, QComboBox, QDialog, QFormLayout,
                              QHBoxLayout, QLabel, QMessageBox, QPushButton,
                              QSlider, QSpinBox, QVBoxLayout, QGroupBox,
@@ -408,7 +406,7 @@ class DeployPredictorDialog(QDialog):
             self.loaded_model.predict(
                 np.zeros((640, 640, 3), dtype=np.uint8),
                 visual_prompts=visuals.copy(),  # This needs to happen to properly initialize the predictor
-                predictor=YOLOEVPSegPredictor,  # This also needs to be SegPredictor, no matter what
+                predictor=YOLOEVPDetectPredictor if self.task == 'detect' else YOLOEVPSegPredictor,
                 imgsz=640,
                 conf=0.99,
             )
@@ -545,12 +543,15 @@ class DeployPredictorDialog(QDialog):
 
         # Get the scaled visual prompts
         visual_prompts = self.scale_prompts(bboxes, masks)
+        
+        # Set the predictor
+        predictor=YOLOEVPDetectPredictor if self.task == 'detect' else YOLOEVPSegPredictor
 
         try:
             # Make predictions
             results = self.loaded_model.predict(self.resized_image,
-                                                visual_prompts=visual_prompts.copy(),  
-                                                predictor=YOLOEVPSegPredictor,
+                                                visual_prompts=visual_prompts.copy(),
+                                                predictor=predictor,
                                                 imgsz=max(self.resized_image.shape[:2]),
                                                 conf=self.main_window.get_uncertainty_thresh(),
                                                 iou=self.main_window.get_iou_thresh(),
@@ -615,6 +616,9 @@ class DeployPredictorDialog(QDialog):
         progress_bar = ProgressBar(self.annotation_window, title="Making Predictions")
         progress_bar.show()
         progress_bar.start_progress(len(target_images))
+        
+        # Set the predictor
+        predictor = YOLOEVPDetectPredictor if self.task == 'detect' else YOLOEVPSegPredictor
 
         for target_image in target_images:
 
@@ -623,7 +627,7 @@ class DeployPredictorDialog(QDialog):
                 results = self.loaded_model.predict(target_image,
                                                     refer_image=refer_image,
                                                     visual_prompts=visual_prompts.copy(),
-                                                    predictor=YOLOEVPSegPredictor,
+                                                    predictor=predictor,
                                                     imgsz=self.imgsz_spinbox.value(),
                                                     conf=self.main_window.get_uncertainty_thresh(),
                                                     iou=self.main_window.get_iou_thresh(),

coralnet_toolbox/Tile/QtTileBatchInference.py

@@ -37,8 +37,8 @@ class TileBatchInference(QDialog):
         self.detect_dialog = main_window.detect_deploy_model_dialog
         self.segment_dialog = main_window.segment_deploy_model_dialog
         self.sam_dialog = main_window.sam_deploy_generator_dialog
-        self.autodistill_dialog = main_window.auto_distill_deploy_model_dialog
-        
+        self.transformers_dialog = main_window.transformers_deploy_model_dialog
+
         # Create a dictionary of the different model dialogs and their loaded models
         self.model_dialogs = {}
 
@@ -167,8 +167,8 @@ class TileBatchInference(QDialog):
             self.model_dialogs["Segment"] = self.segment_dialog
         if self.sam_dialog and getattr(self.sam_dialog, "loaded_model", None):
             self.model_dialogs["SAM Generator"] = self.sam_dialog
-        if self.autodistill_dialog and getattr(self.autodistill_dialog, "loaded_model", None):
-            self.model_dialogs["Autodistill"] = self.autodistill_dialog
+        if self.transformers_dialog and getattr(self.transformers_dialog, "loaded_model", None):
+            self.model_dialogs["Transformers"] = self.transformers_dialog
             
         # Update the model combo box with the available models
         self.update_model_combo()