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

coralnet_toolbox/SeeAnything/QtDeployPredictor.py

@@ -2,10 +2,10 @@ import warnings
 
 import os
 import gc
-import ujson as json
 
 import numpy as np
 
+import torch
 from torch.cuda import empty_cache
 from ultralytics.utils import ops
 
@@ -17,7 +17,7 @@ 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, QTabWidget,
+                             QSlider, QSpinBox, QVBoxLayout, QGroupBox,
                              QWidget, QLineEdit, QFileDialog)
 
 from coralnet_toolbox.Results import ResultsProcessor
@@ -98,7 +98,10 @@ class DeployPredictorDialog(QDialog):
         layout = QVBoxLayout()
 
         # Create a QLabel with explanatory text and hyperlink
-        info_label = QLabel("Choose a Predictor to deploy and use interactively with the See Anything tool.")
+        info_label = QLabel(
+            "Choose a Predictor to deploy and use interactively with the See Anything tool. "
+            "Optionally include a custom visual prompt encoding (VPE) file."
+        )
 
         info_label.setOpenExternalLinks(True)
         info_label.setWordWrap(True)
@@ -109,21 +112,15 @@ class DeployPredictorDialog(QDialog):
 
     def setup_models_layout(self):
         """
-        Setup the models layout with tabs for standard and custom models.
+        Setup the models layout with standard models and file selection.
         """
         group_box = QGroupBox("Model Selection")
-        layout = QVBoxLayout()
-
-        # Create tabbed widget
-        tab_widget = QTabWidget()
-
-        # Tab 1: Standard models
-        standard_tab = QWidget()
-        standard_layout = QVBoxLayout(standard_tab)
-
+        layout = QFormLayout()
+    
+        # Model dropdown
         self.model_combo = QComboBox()
         self.model_combo.setEditable(True)
-
+    
         # Define available models
         standard_models = [
             'yoloe-v8s-seg.pt',
@@ -133,49 +130,15 @@ class DeployPredictorDialog(QDialog):
             'yoloe-11m-seg.pt',
             'yoloe-11l-seg.pt',
         ]
-
+    
         # Add all models to combo box
         self.model_combo.addItems(standard_models)
+        
         # Set the default model
         self.model_combo.setCurrentIndex(standard_models.index('yoloe-v8s-seg.pt'))
-
-        standard_layout.addWidget(QLabel("Models"))
-        standard_layout.addWidget(self.model_combo)
-
-        tab_widget.addTab(standard_tab, "Use Existing Model")
-
-        # Tab 2: Custom model
-        custom_tab = QWidget()
-        custom_layout = QFormLayout(custom_tab)
-
-        # Custom model file selection
-        self.model_path_edit = QLineEdit()
-        browse_button = QPushButton("Browse...")
-        browse_button.clicked.connect(self.browse_model_file)
-
-        model_path_layout = QHBoxLayout()
-        model_path_layout.addWidget(self.model_path_edit)
-        model_path_layout.addWidget(browse_button)
-        custom_layout.addRow("Custom Model:", model_path_layout)
-
-        # Class Mapping
-        self.mapping_edit = QLineEdit()
-        self.mapping_button = QPushButton("Browse...")
-        self.mapping_button.clicked.connect(self.browse_class_mapping_file)
-
-        class_mapping_layout = QHBoxLayout()
-        class_mapping_layout.addWidget(self.mapping_edit)
-        class_mapping_layout.addWidget(self.mapping_button)
-        custom_layout.addRow("Class Mapping:", class_mapping_layout)
-
-        tab_widget.addTab(custom_tab, "Custom Model")
-
-        # Add the tab widget to the main layout
-        layout.addWidget(tab_widget)
-
-        # Store the tab widget for later reference
-        self.model_tab_widget = tab_widget
-
+        # Create a layout for the model selection
+        layout.addRow("Models:", self.model_combo)
+        
         group_box.setLayout(layout)
         self.layout.addWidget(group_box)
 
@@ -208,7 +171,7 @@ class DeployPredictorDialog(QDialog):
         # Image size control
         self.imgsz_spinbox = QSpinBox()
         self.imgsz_spinbox.setRange(512, 65536)
-        self.imgsz_spinbox.setSingleStep(256)
+        self.imgsz_spinbox.setSingleStep(1024)
         self.imgsz_spinbox.setValue(self.imgsz)
         layout.addRow("Image Size (imgsz)", self.imgsz_spinbox)
 
@@ -304,37 +267,6 @@ class DeployPredictorDialog(QDialog):
 
         group_box.setLayout(layout)
         self.layout.addWidget(group_box)
-        
-    def browse_model_file(self):
-        """
-        Open a file dialog to browse for a model file.
-        """
-        file_path, _ = QFileDialog.getOpenFileName(self,
-                                                   "Select Model File",
-                                                   "",
-                                                   "Model Files (*.pt *.pth);;All Files (*)")
-        if file_path:
-            self.model_path_edit.setText(file_path)
-
-            # Load the class mapping if it exists
-            dir_path = os.path.dirname(os.path.dirname(file_path))
-            class_mapping_path = f"{dir_path}/class_mapping.json"
-            if os.path.exists(class_mapping_path):
-                self.class_mapping = json.load(open(class_mapping_path, 'r'))
-                self.mapping_edit.setText(class_mapping_path)
-
-    def browse_class_mapping_file(self):
-        """
-        Browse and select a class mapping file.
-        """
-        file_path, _ = QFileDialog.getOpenFileName(self,
-                                                   "Select Class Mapping File",
-                                                   "",
-                                                   "JSON Files (*.json)")
-        if file_path:
-            # Load the class mapping
-            self.class_mapping = json.load(open(file_path, 'r'))
-            self.mapping_edit.setText(file_path)
 
     def initialize_uncertainty_threshold(self):
         """Initialize the uncertainty threshold slider with the current value"""
@@ -390,6 +322,7 @@ class DeployPredictorDialog(QDialog):
     def is_sam_model_deployed(self):
         """
         Check if the SAM model is deployed and update the checkbox state accordingly.
+        If SAM is enabled for polygons, sync and disable the imgsz spinbox.
 
         :return: Boolean indicating whether the SAM model is deployed
         """
@@ -402,9 +335,48 @@ class DeployPredictorDialog(QDialog):
             self.use_sam_dropdown.setCurrentText("False")
             QMessageBox.critical(self, "Error", "Please deploy the SAM model first.")
             return False
+        
+        # Check if SAM polygons are enabled
+        if self.use_sam_dropdown.currentText() == "True":
+            # Sync the imgsz spinbox with SAM's value
+            self.imgsz_spinbox.setValue(self.sam_dialog.imgsz_spinbox.value())
+            # Disable the spinbox
+            self.imgsz_spinbox.setEnabled(False)
+            
+            # Connect SAM's imgsz_spinbox valueChanged signal to update our value
+            # First disconnect any existing connection to avoid duplicates
+            try:
+                self.sam_dialog.imgsz_spinbox.valueChanged.disconnect(self.update_from_sam_imgsz)
+            except TypeError:
+                # No connection exists yet
+                pass
+            
+            # Connect the signal
+            self.sam_dialog.imgsz_spinbox.valueChanged.connect(self.update_from_sam_imgsz)
+        else:
+            # Re-enable the spinbox when SAM polygons are disabled
+            self.imgsz_spinbox.setEnabled(True)
+            
+            # Disconnect the signal when SAM is disabled
+            try:
+                self.sam_dialog.imgsz_spinbox.valueChanged.disconnect(self.update_from_sam_imgsz)
+            except TypeError:
+                # No connection exists
+                pass
 
         return True
 
+    def update_from_sam_imgsz(self, value):
+        """
+        Update the SeeAnything image size when SAM's image size changes.
+        Only takes effect when SAM polygons are enabled.
+        
+        Args:
+            value (int): The new image size value from SAM dialog
+        """
+        if self.use_sam_dropdown.currentText() == "True":
+            self.imgsz_spinbox.setValue(value)
+
     def load_model(self):
         """
         Load the selected model.
@@ -412,46 +384,45 @@ class DeployPredictorDialog(QDialog):
         QApplication.setOverrideCursor(Qt.WaitCursor)
         progress_bar = ProgressBar(self.annotation_window, title="Loading Model")
         progress_bar.show()
-
+    
         try:
             # Get selected model path and download weights if needed
             self.model_path = self.model_combo.currentText()
-
+    
             # Load model using registry
             self.loaded_model = YOLOE(self.model_path).to(self.main_window.device)
-
-            # Create a dummy visual dictionary
+    
+            # Create a dummy visual dictionary for standard model loading
             visuals = dict(
                 bboxes=np.array(
                     [
-                        [120, 425, 160, 445],
+                        [120, 425, 160, 445],  # Random box
                     ],
                 ),
                 cls=np.array(
                     np.zeros(1),
                 ),
             )
-
+    
             # Run a dummy prediction to load the model
             self.loaded_model.predict(
                 np.zeros((640, 640, 3), dtype=np.uint8),
-                visual_prompts=visuals.copy(),
-                predictor=YOLOEVPDetectPredictor,
+                visual_prompts=visuals.copy(),  # This needs to happen to properly initialize the predictor
+                predictor=YOLOEVPSegPredictor,  # This also needs to be SegPredictor, no matter what
                 imgsz=640,
                 conf=0.99,
             )
-
-            # Load the model class names if available
-            if self.class_mapping:
-                self.add_labels_to_label_window()
-
+            # Finish the progress bar
             progress_bar.finish_progress()
-            self.status_bar.setText("Model loaded")
+            # Update the status bar
+            self.status_bar.setText(f"Loaded ({self.model_path}")
             QMessageBox.information(self.annotation_window, "Model Loaded", "Model loaded successfully")
 
         except Exception as e:
+            self.loaded_model = None
+            self.status_bar.setText(f"Error loading model: {self.model_path}")
             QMessageBox.critical(self.annotation_window, "Error Loading Model", f"Error loading model: {e}")
-
+    
         finally:
             # Restore cursor
             QApplication.restoreOverrideCursor()
@@ -459,19 +430,7 @@ class DeployPredictorDialog(QDialog):
             progress_bar.stop_progress()
             progress_bar.close()
             progress_bar = None
-
-        self.accept()
-
-    def add_labels_to_label_window(self):
-        """
-        Add labels to the label window based on the class mapping.
-        """
-        if self.class_mapping:
-            for label in self.class_mapping.values():
-                self.main_window.label_window.add_label_if_not_exists(label['short_label_code'],
-                                                                      label['long_label_code'],
-                                                                      QColor(*label['color']))
-
+            
     def resize_image(self, image):
         """
         Resize the image to the specified size.
@@ -526,26 +485,11 @@ class DeployPredictorDialog(QDialog):
             self.resized_image = self.resize_image(image)
         else:
             self.resized_image = image
-
-    def predict_from_prompts(self, bboxes, masks=None):
+            
+    def scale_prompts(self, bboxes, masks=None):
         """
-        Make predictions using the currently loaded model using prompts.
-
-        Args:
-            bboxes (np.ndarray): The bounding boxes to use as prompts.
-            masks (list, optional): A list of polygons to use as prompts for segmentation.
-
-        Returns:
-            results (Results): Ultralytics Results object
+        Scale the bounding boxes and masks to the resized image.
         """
-        if not self.loaded_model:
-            QMessageBox.critical(self.annotation_window, "Model Not Loaded",
-                                 "Model not loaded, cannot make predictions")
-            return None
-
-        if not len(bboxes):
-            return None
-
         # Update the bbox coordinates to be relative to the resized image
         bboxes = np.array(bboxes)
         bboxes[:, 0] = (bboxes[:, 0] / self.original_image.shape[1]) * self.resized_image.shape[1]
@@ -557,7 +501,7 @@ class DeployPredictorDialog(QDialog):
         self.task = self.task_dropdown.currentText()
 
         # Create a visual dictionary
-        visuals = {
+        visual_prompts = {
             'bboxes': np.array(bboxes),
             'cls': np.zeros(len(bboxes))
         }
@@ -569,21 +513,44 @@ class DeployPredictorDialog(QDialog):
                     scaled_mask[:, 0] = (scaled_mask[:, 0] / self.original_image.shape[1]) * self.resized_image.shape[1]
                     scaled_mask[:, 1] = (scaled_mask[:, 1] / self.original_image.shape[0]) * self.resized_image.shape[0]
                     scaled_masks.append(scaled_mask)
-                visuals['masks'] = scaled_masks
+                visual_prompts['masks'] = scaled_masks
             else:  # Fallback to creating masks from bboxes if no masks are provided
                 fallback_masks = []
                 for bbox in bboxes:
                     x1, y1, x2, y2 = bbox
                     fallback_masks.append(np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]))
-                visuals['masks'] = fallback_masks
+                visual_prompts['masks'] = fallback_masks
+        
+        return visual_prompts
 
-        predictor = YOLOEVPSegPredictor if self.task == "segment" else YOLOEVPDetectPredictor
+    def predict_from_prompts(self, bboxes, masks=None):
+        """
+        Make predictions using the currently loaded model using prompts.
+
+        Args:
+            bboxes (np.ndarray): The bounding boxes to use as prompts.
+            masks (list, optional): A list of polygons to use as prompts for segmentation.
+
+        Returns:
+            results (Results): Ultralytics Results object
+        """
+        if not self.loaded_model:
+            QMessageBox.critical(self.annotation_window, 
+                                 "Model Not Loaded",
+                                 "Model not loaded, cannot make predictions")
+            return None
+
+        if not len(bboxes):
+            return None
+
+        # Get the scaled visual prompts
+        visual_prompts = self.scale_prompts(bboxes, masks)
 
         try:
             # Make predictions
             results = self.loaded_model.predict(self.resized_image,
-                                                visual_prompts=visuals.copy(),
-                                                predictor=predictor,
+                                                visual_prompts=visual_prompts.copy(),  
+                                                predictor=YOLOEVPSegPredictor,
                                                 imgsz=max(self.resized_image.shape[:2]),
                                                 conf=self.main_window.get_uncertainty_thresh(),
                                                 iou=self.main_window.get_iou_thresh(),
@@ -618,18 +585,30 @@ class DeployPredictorDialog(QDialog):
             max_area_thresh=self.main_window.get_area_thresh_max()
         )
 
-        # Set the predictor
-        self.task = self.task_dropdown.currentText()
-
-        # Create a visual dictionary
-        visuals = {
-            'bboxes': np.array(refer_bboxes),
-            'cls': np.zeros(len(refer_bboxes))
-        }
-        if self.task == 'segment':
-            visuals['masks'] = refer_masks
-
-        predictor = YOLOEVPSegPredictor if self.task == "segment" else YOLOEVPDetectPredictor
+        # Get the scaled visual prompts
+        visual_prompts = self.scale_prompts(refer_bboxes, refer_masks)
+
+        # If VPEs are being used
+        if self.vpe is not None:
+            # Generate a new VPE from the current visual prompts
+            new_vpe = self.prompts_to_vpes(visual_prompts, self.resized_image)
+            
+            if new_vpe is not None:
+                # If we already have a VPE, average with the existing one
+                if self.vpe.shape == new_vpe.shape:
+                    self.vpe = (self.vpe + new_vpe) / 2
+                    # Re-normalize
+                    self.vpe = torch.nn.functional.normalize(self.vpe, p=2, dim=-1)
+                else:
+                    # Replace with the new VPE if shapes don't match
+                    self.vpe = new_vpe
+                
+                # Set the updated VPE in the model
+                self.loaded_model.is_fused = lambda: False
+                self.loaded_model.set_classes(["object0"], self.vpe)
+            
+            # Clear visual prompts since we're using VPE
+            visual_prompts = {}  # this is okay with a fused model
 
         # Create a progress bar
         QApplication.setOverrideCursor(Qt.WaitCursor)
@@ -643,8 +622,8 @@ class DeployPredictorDialog(QDialog):
                 # Make predictions
                 results = self.loaded_model.predict(target_image,
                                                     refer_image=refer_image,
-                                                    visual_prompts=visuals.copy(),
-                                                    predictor=predictor,
+                                                    visual_prompts=visual_prompts.copy(),
+                                                    predictor=YOLOEVPSegPredictor,
                                                     imgsz=self.imgsz_spinbox.value(),
                                                     conf=self.main_window.get_uncertainty_thresh(),
                                                     iou=self.main_window.get_iou_thresh(),

coralnet_toolbox/Tools/QtCutSubTool.py

@@ -87,9 +87,16 @@ class CutSubTool(SubTool):
         self._update_cut_line_path(position)
 
     def keyPressEvent(self, event):
-        """Handle key press events for canceling the cut."""
-        if event.key() in (Qt.Key_Backspace, Qt.Key_Escape):
+        """Handle key press events for cutting operations."""
+        # Check for Ctrl+X to toggle cutting mode off
+        if event.modifiers() & Qt.ControlModifier and event.key() == Qt.Key_X:
             self.parent_tool.deactivate_subtool()
+            return
+            
+        # Handle Backspace to clear the current cutting line but stay in cutting mode
+        if event.key() == Qt.Key_Backspace:
+            self._clear_cutting_line()
+            return
 
     def _start_drawing_cut_line(self, position):
         """Start drawing the cut line from the given position."""
@@ -115,6 +122,15 @@ class CutSubTool(SubTool):
                 path.lineTo(point)
             self.cutting_path_item.setPath(path)
             
+    def _clear_cutting_line(self):
+        """Clear the current cutting line but remain in cutting mode."""
+        self.cutting_points = []
+        self.drawing_in_progress = False
+        if self.cutting_path_item:
+            self.annotation_window.scene.removeItem(self.cutting_path_item)
+            self.cutting_path_item = None
+            self.annotation_window.scene.update()
+            
     def _break_apart_multipolygon(self):
         """Handle the special case of 'cutting' a MultiPolygonAnnotation."""
         new_annotations = self.target_annotation.cut()

coralnet_toolbox/Tools/QtPolygonTool.py

@@ -133,12 +133,51 @@ class PolygonTool(Tool):
             return None
 
         # Create the annotation with current points
-        # The polygon simplification is now handled inside the PolygonAnnotation class
         if finished and len(self.points) > 2:
             # Close the polygon
             self.points.append(self.points[0])
-
-        # Create the annotation - will be simplified in the constructor
+            
+            # --- Validation for polygon size and shape ---
+            # Step 1: Remove duplicate or near-duplicate points
+            filtered_points = []
+            MIN_DISTANCE = 2.0  # Minimum distance between points in pixels
+            
+            for i, point in enumerate(self.points):
+                # Skip if this point is too close to the previous one
+                if i > 0:
+                    prev_point = filtered_points[-1]
+                    distance = ((point.x() - prev_point.x())**2 + (point.y() - prev_point.y())**2)**0.5
+                    if distance < MIN_DISTANCE:
+                        continue
+                filtered_points.append(point)
+            
+            # Step 2: Ensure we have enough points for a valid polygon
+            if len(filtered_points) < 4:  # Need at least 3 + 1 closing point
+                # Create a small triangle/square if we don't have enough points
+                if len(filtered_points) > 0:
+                    center_x = sum(p.x() for p in filtered_points) / len(filtered_points)
+                    center_y = sum(p.y() for p in filtered_points) / len(filtered_points)
+                    
+                    # Create a small polygon centered on the average of existing points
+                    MIN_SIZE = 5.0
+                    filtered_points = [
+                        QPointF(center_x - MIN_SIZE, center_y - MIN_SIZE),
+                        QPointF(center_x + MIN_SIZE, center_y - MIN_SIZE),
+                        QPointF(center_x + MIN_SIZE, center_y + MIN_SIZE),
+                        QPointF(center_x - MIN_SIZE, center_y + MIN_SIZE),
+                        QPointF(center_x - MIN_SIZE, center_y - MIN_SIZE)  # Close the polygon
+                    ]
+                    
+                    QMessageBox.information(
+                        self.annotation_window, 
+                        "Polygon Adjusted",
+                        "The polygon had too few unique points and has been adjusted to a minimum size."
+                    )
+            
+            # Use the filtered points list instead of the original
+            self.points = filtered_points
+
+        # Create the annotation with validated points
         annotation = PolygonAnnotation(self.points,
                                        self.annotation_window.selected_label.short_label_code,
                                        self.annotation_window.selected_label.long_label_code,

coralnet_toolbox/Tools/QtRectangleTool.py

@@ -113,6 +113,36 @@ class RectangleTool(Tool):
         # Ensure top_left and bottom_right are correctly calculated
         top_left = QPointF(min(self.start_point.x(), end_point.x()), min(self.start_point.y(), end_point.y()))
         bottom_right = QPointF(max(self.start_point.x(), end_point.x()), max(self.start_point.y(), end_point.y()))
+        
+        # Calculate width and height of the rectangle
+        width = bottom_right.x() - top_left.x()
+        height = bottom_right.y() - top_left.y()
+        
+        # Define minimum dimensions for a valid rectangle (e.g., 3x3 pixels)
+        MIN_DIMENSION = 3.0
+        
+        # If rectangle is too small and we're finalizing it, enforce minimum size
+        if finished and (width < MIN_DIMENSION or height < MIN_DIMENSION):
+            if width < MIN_DIMENSION:
+                # Expand width while maintaining center
+                center_x = (top_left.x() + bottom_right.x()) / 2
+                top_left.setX(center_x - MIN_DIMENSION / 2)
+                bottom_right.setX(center_x + MIN_DIMENSION / 2)
+                
+            if height < MIN_DIMENSION:
+                # Expand height while maintaining center
+                center_y = (top_left.y() + bottom_right.y()) / 2
+                top_left.setY(center_y - MIN_DIMENSION / 2)
+                bottom_right.setY(center_y + MIN_DIMENSION / 2)
+            
+            # Show a message if we had to adjust a very small rectangle
+            if width < 1 or height < 1:
+                QMessageBox.information(
+                    self.annotation_window, 
+                    "Rectangle Adjusted",
+                    f"The rectangle was too small and has been adjusted to a minimum size of "
+                    f"{MIN_DIMENSION}x{MIN_DIMENSION} pixels."
+                )
 
         # Create the rectangle annotation
         annotation = RectangleAnnotation(top_left,

coralnet_toolbox/Tools/QtResizeSubTool.py

@@ -124,5 +124,22 @@ class ResizeSubTool(SubTool):
         }
 
     def _get_polygon_handles(self, annotation):
-        """Return resize handles for a polygon annotation."""
-        return {f"point_{i}": QPointF(p.x(), p.y()) for i, p in enumerate(annotation.points)}
+        """
+        Return resize handles for a polygon, including its outer boundary and all holes.
+        Uses the new handle format: 'point_{poly_index}_{vertex_index}'.
+        """
+        handles = {}
+
+        # 1. Create handles for the outer boundary using the 'outer' keyword.
+        for i, p in enumerate(annotation.points):
+            handle_name = f"point_outer_{i}"
+            handles[handle_name] = QPointF(p.x(), p.y())
+
+        # 2. Create handles for each of the inner holes using their index.
+        if hasattr(annotation, 'holes'):
+            for hole_index, hole in enumerate(annotation.holes):
+                for vertex_index, p in enumerate(hole):
+                    handle_name = f"point_{hole_index}_{vertex_index}"
+                    handles[handle_name] = QPointF(p.x(), p.y())
+        
+        return handles