coralnet-toolbox 0.0.73__py2.py3-none-any.whl → 0.0.74__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)
 
@@ -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"""
@@ -412,46 +344,43 @@ 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()
-
-            progress_bar.finish_progress()
-            self.status_bar.setText("Model loaded")
+            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()
@@ -460,18 +389,6 @@ class DeployPredictorDialog(QDialog):
             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 +443,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 +459,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 +471,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
+
+    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
 
-        predictor = YOLOEVPSegPredictor if self.task == "segment" else YOLOEVPDetectPredictor
+        # 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 +543,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 +580,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/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

coralnet_toolbox/Tools/QtSAMTool.py

@@ -1,7 +1,7 @@
 import warnings
 import numpy as np
 
-from PyQt5.QtCore import Qt, QPointF, QRectF, QTimer
+from PyQt5.QtCore import Qt, QPointF, QRectF
 from PyQt5.QtGui import QMouseEvent, QKeyEvent, QPen, QColor, QBrush, QPainterPath
 from PyQt5.QtWidgets import QMessageBox, QGraphicsEllipseItem, QGraphicsRectItem, QGraphicsPathItem, QApplication
 
@@ -12,6 +12,7 @@ from coralnet_toolbox.QtWorkArea import WorkArea
 
 from coralnet_toolbox.utilities import pixmap_to_numpy
 from coralnet_toolbox.utilities import simplify_polygon
+from coralnet_toolbox.utilities import polygonize_mask_with_holes
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -369,40 +370,47 @@ class SAMTool(Tool):
                 QApplication.restoreOverrideCursor()
                 return
 
-            # Get the points of the top1 mask
+            # Get the top confidence prediction's mask tensor
             top1_index = np.argmax(results.boxes.conf)
-            predictions = results[top1_index].masks.xy[0]
+            mask_tensor = results[top1_index].masks.data
 
-            # Safety check: make sure we have predicted points
-            if len(predictions) == 0:
-                QApplication.restoreOverrideCursor()
-                return
+            # Check if holes are allowed from the SAM dialog
+            allow_holes = self.sam_dialog.get_allow_holes()
 
-            # Clean the polygon using Ramer-Douglas-Peucker algorithm
-            predictions = simplify_polygon(predictions, 0.1)
+            # Polygonize the mask to get the exterior and holes
+            exterior_coords, holes_coords_list = polygonize_mask_with_holes(mask_tensor)
 
             # Safety check: need at least 3 points for a valid polygon
-            if len(predictions) < 3:
+            if len(exterior_coords) < 3:
                 QApplication.restoreOverrideCursor()
                 return
 
-            # Move the points back to the original image space
+            # --- Process and Clean the Polygon Points ---
             working_area_top_left = self.working_area.rect.topLeft()
-            points = [(point[0] + working_area_top_left.x(),
-                       point[1] + working_area_top_left.y()) for point in predictions]
+            offset_x, offset_y = working_area_top_left.x(), working_area_top_left.y()
 
-            # Convert to QPointF for graphics
-            self.points = [QPointF(*point) for point in points]
+            # Simplify, offset, and convert the exterior points
+            simplified_exterior = simplify_polygon(exterior_coords, 0.1)
+            self.points = [QPointF(p[0] + offset_x, p[1] + offset_y) for p in simplified_exterior]
 
-            # Create the temporary annotation
+            # Simplify, offset, and convert each hole only if allowed
+            final_holes = []
+            if allow_holes:
+                for hole_coords in holes_coords_list:
+                    if len(hole_coords) >= 3:  # Ensure holes are also valid polygons
+                        simplified_hole = simplify_polygon(hole_coords, 0.1)
+                        final_holes.append([QPointF(p[0] + offset_x, p[1] + offset_y) for p in simplified_hole])
+
+            # Create the temporary annotation, now with holes (or not)
             self.temp_annotation = PolygonAnnotation(
-                self.points,
-                self.annotation_window.selected_label.short_label_code,
-                self.annotation_window.selected_label.long_label_code,
-                self.annotation_window.selected_label.color,
-                self.annotation_window.current_image_path,
-                self.annotation_window.selected_label.id,
-                self.main_window.label_window.active_label.transparency
+                points=self.points,
+                holes=final_holes,
+                short_label_code=self.annotation_window.selected_label.short_label_code,
+                long_label_code=self.annotation_window.selected_label.long_label_code,
+                color=self.annotation_window.selected_label.color,
+                image_path=self.annotation_window.current_image_path,
+                label_id=self.annotation_window.selected_label.id,
+                transparency=self.main_window.label_window.active_label.transparency
             )
 
             # Create the graphics item for the temporary annotation
@@ -611,12 +619,13 @@ class SAMTool(Tool):
                     # Use existing temporary annotation
                     final_annotation = PolygonAnnotation(
                         self.points,
-                        self.annotation_window.selected_label.short_label_code,
-                        self.annotation_window.selected_label.long_label_code,
-                        self.annotation_window.selected_label.color,
-                        self.annotation_window.current_image_path,
-                        self.annotation_window.selected_label.id,
-                        self.main_window.label_window.active_label.transparency
+                        self.temp_annotation.label.short_label_code,
+                        self.temp_annotation.label.long_label_code,
+                        self.temp_annotation.label.color,
+                        self.temp_annotation.image_path,
+                        self.temp_annotation.label.id,
+                        self.temp_annotation.label.transparency,
+                        holes=self.temp_annotation.holes
                     )
 
                     # Copy confidence data
@@ -637,7 +646,7 @@ class SAMTool(Tool):
                     final_annotation = self.create_annotation(True)
                     if final_annotation:
                         self.annotation_window.add_annotation_from_tool(final_annotation)
-                        self.clear_prompt_graphics()
+                        self.clear_prompt_graphics() 
             # If no active prompts, cancel the working area
             else:
                 self.cancel_working_area()
@@ -727,24 +736,36 @@ class SAMTool(Tool):
             QApplication.restoreOverrideCursor()
             return None
 
-        # Get the top confidence prediction
+        # Get the top confidence prediction's mask tensor
         top1_index = np.argmax(results.boxes.conf)
-        predictions = results[top1_index].masks.xy[0]
+        mask_tensor = results[top1_index].masks.data
+
+        # Check if holes are allowed from the SAM dialog
+        allow_holes = self.sam_dialog.get_allow_holes()
 
-        # Safety check for predictions
-        if len(predictions) == 0:
+        # Polygonize the mask using the new method to get the exterior and holes
+        exterior_coords, holes_coords_list = polygonize_mask_with_holes(mask_tensor)
+
+        # Safety check for an empty result
+        if not exterior_coords:
             QApplication.restoreOverrideCursor()
             return None
 
-        # Clean polygon points
-        predictions = simplify_polygon(predictions, 0.1)
-
-        # Move points back to original image space
+        # --- Process and Clean the Polygon Points ---
         working_area_top_left = self.working_area.rect.topLeft()
-        points = [(point[0] + working_area_top_left.x(),
-                   point[1] + working_area_top_left.y()) for point in predictions]
-        # Convert to QPointF for graphics
-        self.points = [QPointF(*point) for point in points]
+        offset_x, offset_y = working_area_top_left.x(), working_area_top_left.y()
+
+        # Simplify, offset, and convert the exterior points
+        simplified_exterior = simplify_polygon(exterior_coords, 0.1)
+        self.points = [QPointF(p[0] + offset_x, p[1] + offset_y) for p in simplified_exterior]
+
+        # Simplify, offset, and convert each hole only if allowed
+        final_holes = []
+        if allow_holes:
+            for hole_coords in holes_coords_list:
+                if len(hole_coords) >= 3:
+                    simplified_hole = simplify_polygon(hole_coords, 0.1)
+                    final_holes.append([QPointF(p[0] + offset_x, p[1] + offset_y) for p in simplified_hole])
 
         # Require at least 3 points for valid polygon
         if len(self.points) < 3:
@@ -754,15 +775,16 @@ class SAMTool(Tool):
         # Get confidence score
         confidence = results.boxes.conf[top1_index].item()
 
-        # Create final annotation
+        # Create final annotation, now passing the holes argument
         annotation = PolygonAnnotation(
-            self.points,
-            self.annotation_window.selected_label.short_label_code,
-            self.annotation_window.selected_label.long_label_code,
-            self.annotation_window.selected_label.color,
-            self.annotation_window.current_image_path,
-            self.annotation_window.selected_label.id,
-            self.main_window.label_window.active_label.transparency
+            points=self.points,
+            holes=final_holes,
+            short_label_code=self.annotation_window.selected_label.short_label_code,
+            long_label_code=self.annotation_window.selected_label.long_label_code,
+            color=self.annotation_window.selected_label.color,
+            image_path=self.annotation_window.current_image_path,
+            label_id=self.annotation_window.selected_label.id,
+            transparency=self.main_window.label_window.active_label.transparency
         )
 
         # Update confidence

coralnet_toolbox/Tools/QtSeeAnythingTool.py

@@ -173,6 +173,7 @@ class SeeAnythingTool(Tool):
 
         # Set the image in the SeeAnything dialog
         self.see_anything_dialog.set_image(self.work_area_image, self.image_path)
+        # self.see_anything_dialog.reload_model()
 
         self.annotation_window.setCursor(Qt.CrossCursor)
         self.annotation_window.scene.update()
@@ -552,9 +553,9 @@ class SeeAnythingTool(Tool):
         # Move the points back to the original image space
         working_area_top_left = self.working_area.rect.topLeft()
 
-        task = self.see_anything_dialog.task_dropdown.currentText()
         masks = None
-        if task == 'segment':
+        # Create masks from the rectangles (these are not polygons)
+        if self.see_anything_dialog.task_dropdown.currentText() == 'segment':
             masks = []
             for r in self.rectangles:
                 x1, y1, x2, y2 = r
@@ -587,8 +588,8 @@ class SeeAnythingTool(Tool):
         # Clear previous annotations if any
         self.clear_annotations()
 
-        # Process results based on the task type
-        if self.see_anything_dialog.task == "segment":
+        # Process results based on the task type (creates polygons or rectangle annotations)
+        if self.see_anything_dialog.task_dropdown.currentText() == "segment":
             if self.results.masks:
                 for i, polygon in enumerate(self.results.masks.xyn):
                     confidence = self.results.boxes.conf[i].item()
@@ -624,7 +625,9 @@ class SeeAnythingTool(Tool):
                     box_abs_work_area = box_norm.detach().cpu().numpy() * np.array(
                         [self.work_area_image.shape[1], self.work_area_image.shape[0],
                          self.work_area_image.shape[1], self.work_area_image.shape[0]])
-                    box_area = (box_abs_work_area[2] - box_abs_work_area[0]) * (box_abs_work_area[3] - box_abs_work_area[1])
+                    # Calculate the area of the bounding box
+                    box_area = (box_abs_work_area[2] - box_abs_work_area[0]) * \
+                               (box_abs_work_area[3] - box_abs_work_area[1])
 
                     # Area filtering
                     min_area = self.main_window.get_area_thresh_min() * image_area