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

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
@@ -416,7 +418,7 @@ class DeployGeneratorDialog(QDialog):
 
         # Image size control
         self.imgsz_spinbox = QSpinBox()
-        self.imgsz_spinbox.setRange(512, 65536)
+        self.imgsz_spinbox.setRange(1024, 65536)
         self.imgsz_spinbox.setSingleStep(1024)
         self.imgsz_spinbox.setValue(self.imgsz)
         layout.addRow("Image Size (imgsz):", self.imgsz_spinbox)
@@ -502,8 +504,13 @@ class DeployGeneratorDialog(QDialog):
 
         main_layout.addLayout(button_row)
 
-        # Second row: Save VPE button + Show VPE button side by side
+        # Second row: VPE action buttons
         vpe_row = QHBoxLayout()
+        
+        generate_vpe_button = QPushButton("Generate VPEs")
+        generate_vpe_button.clicked.connect(self.generate_vpes_from_references)
+        vpe_row.addWidget(generate_vpe_button)
+
         save_vpe_button = QPushButton("Save VPE")
         save_vpe_button.clicked.connect(self.save_vpe)
         vpe_row.addWidget(save_vpe_button)
@@ -762,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)")
@@ -810,80 +817,75 @@ class DeployGeneratorDialog(QDialog):
             
     def save_vpe(self):
         """
-        Save the combined collection of VPEs (imported and reference-generated) to disk.
+        Saves the combined collection of VPEs (imported and pre-generated from references) to disk.
         """
-        # Always sync with the live UI selection before saving.
-        self.update_stashed_references_from_ui()
-
-        # Create a list to hold all VPEs
-        all_vpes = []
-        
-        # Add imported VPEs if available
-        if self.imported_vpes:
-            all_vpes.extend(self.imported_vpes)
+        QApplication.setOverrideCursor(Qt.WaitCursor)
         
-        # Check if we should generate new VPEs from reference images
-        references_dict = self._get_references()
-        if references_dict:
-            # Reload the model to ensure clean state
-            self.reload_model()
+        try:
+            # Create a list to hold all VPEs to be saved
+            all_vpes = []
             
-            # Convert references to VPEs without updating self.reference_vpes yet
-            new_vpes = self.references_to_vpe(references_dict, update_reference_vpes=False)
+            # Add imported VPEs if available
+            if self.imported_vpes:
+                all_vpes.extend(self.imported_vpes)
             
-            if new_vpes:
-                # Add new VPEs to collection
-                all_vpes.extend(new_vpes)
-                # Update reference_vpes with the new ones
-                self.reference_vpes = new_vpes
-        else:
-            # Include existing reference VPEs if we have them
+            # Add pre-generated reference VPEs if available
             if self.reference_vpes:
                 all_vpes.extend(self.reference_vpes)
-        
-        # Check if we have any VPEs to save
-        if not all_vpes:
-            QMessageBox.warning(
+            
+            # Check if we have any VPEs to save
+            if not all_vpes:
+                QApplication.restoreOverrideCursor()
+                QMessageBox.warning(
+                    self,
+                    "No VPEs Available",
+                    "No VPEs available to save. "
+                    "Please either load a VPE file or generate VPEs from reference images first."
+                )
+                return
+            
+            QApplication.restoreOverrideCursor()
+            
+            file_path, _ = QFileDialog.getSaveFileName(
                 self,
-                "No VPEs Available",
-                "No VPEs available to save. Please either load a VPE file or select reference images."
+                "Save VPE Collection",
+                "",
+                "PyTorch Tensor (*.pt);;All Files (*)"
             )
-            return
-        
-        # Open file dialog to select save location
-        file_path, _ = QFileDialog.getSaveFileName(
-            self,
-            "Save VPE Collection",
-            "",
-            "PyTorch Tensor (*.pt);;All Files (*)"
-        )
-        
-        if not file_path:
-            return  # User canceled the dialog
-        
-        # Add .pt extension if not present
-        if not file_path.endswith('.pt'):
-            file_path += '.pt'
-        
-        try:
-            # Move tensors to CPU before saving
+            
+            if not file_path:
+                return
+            
+            QApplication.setOverrideCursor(Qt.WaitCursor)
+            
+            if not file_path.endswith('.pt'):
+                file_path += '.pt'
+            
             vpe_list_cpu = [vpe.cpu() for vpe in all_vpes]
             
-            # Save the list of tensors
             torch.save(vpe_list_cpu, file_path)
             
             self.status_bar.setText(f"Saved {len(all_vpes)} VPE tensors to {os.path.basename(file_path)}")
+            
+            QApplication.restoreOverrideCursor()
             QMessageBox.information(
                 self,
                 "VPE Saved",
                 f"Saved {len(all_vpes)} VPE tensors to {file_path}"
             )
+            
         except Exception as e:
+            QApplication.restoreOverrideCursor()
             QMessageBox.critical(
                 self,
                 "Error Saving VPE",
                 f"Failed to save VPE: {str(e)}"
             )
+        finally:
+            try:
+                QApplication.restoreOverrideCursor()
+            except:
+                pass
 
     def load_model(self):
         """
@@ -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(),
@@ -1156,6 +1161,59 @@ class DeployGeneratorDialog(QDialog):
         
         return [[combined_results]]
     
+    def generate_vpes_from_references(self):
+        """
+        Calculates VPEs from the currently highlighted reference images and
+        stores them in self.reference_vpes, overwriting any previous ones.
+        """
+        if not self.loaded_model:
+            QMessageBox.warning(self, "No Model Loaded", "A model must be loaded before generating VPEs.")
+            return
+
+        # Always sync with the live UI selection before generating.
+        self.update_stashed_references_from_ui()
+        references_dict = self._get_references()
+
+        if not references_dict:
+            QMessageBox.information(
+                self,
+                "No References Selected",
+                "Please highlight one or more reference images in the table to generate VPEs."
+            )
+            return
+
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        progress_bar = ProgressBar(self, title="Generating VPEs")
+        progress_bar.show()
+
+        try:
+            # Make progress bar busy
+            progress_bar.set_busy_mode("Generating VPEs...")
+            # Reload the model to ensure a clean state for VPE generation
+            self.reload_model()
+
+            # The references_to_vpe method will calculate and update self.reference_vpes
+            new_vpes = self.references_to_vpe(references_dict, update_reference_vpes=True)
+
+            if new_vpes:
+                num_vpes = len(new_vpes)
+                num_images = len(references_dict)
+                message = f"Successfully generated {num_vpes} VPEs from {num_images} reference image(s)."
+                self.status_bar.setText(message)
+                QMessageBox.information(self, "VPEs Generated", message)
+            else:
+                message = "Could not generate VPEs. Ensure annotations are valid."
+                self.status_bar.setText(message)
+                QMessageBox.warning(self, "Generation Failed", message)
+
+        except Exception as e:
+            QMessageBox.critical(self, "Error Generating VPEs", f"An unexpected error occurred: {str(e)}")
+            self.status_bar.setText("Error during VPE generation.")
+        finally:
+            QApplication.restoreOverrideCursor()
+            progress_bar.stop_progress()
+            progress_bar.close()
+    
     def references_to_vpe(self, reference_dict, update_reference_vpes=True):
         """
         Converts the contents of a reference dictionary to VPEs (Visual Prompt Embeddings).
@@ -1197,14 +1255,13 @@ class DeployGeneratorDialog(QDialog):
             
         return vpe_list
 
-    def _apply_model_using_vpe(self, inputs, references_dict):
+    def _apply_model_using_vpe(self, inputs):
         """
-        Apply the model to the inputs using combined VPEs from both imported files
-        and reference annotations.
+        Apply the model to the inputs using pre-calculated VPEs from imported files
+        and/or generated from reference annotations.
         
         Args:
             inputs (list): List of input images.
-            references_dict (dict): Dictionary containing reference annotations for each image.
             
         Returns:
             list: List of prediction results.
@@ -1219,23 +1276,17 @@ class DeployGeneratorDialog(QDialog):
         if self.imported_vpes:
             combined_vpes.extend(self.imported_vpes)
             
-        # Process reference images to VPEs if any exist
-        if references_dict:
-            # Only update reference_vpes if references_dict is not empty
-            reference_vpes = self.references_to_vpe(references_dict, update_reference_vpes=True)
-            if reference_vpes:
-                combined_vpes.extend(reference_vpes)
-        else:
-            # Use existing reference_vpes if we have them
-            if self.reference_vpes:
-                combined_vpes.extend(self.reference_vpes)
+        # Add pre-generated reference VPEs if available
+        if self.reference_vpes:
+            combined_vpes.extend(self.reference_vpes)
         
         # Check if we have any VPEs to use
         if not combined_vpes:
             QMessageBox.warning(
                 self,
                 "No VPEs Available",
-                "No VPEs available for prediction. Please either load a VPE file or select reference images."
+                "No VPEs are available for prediction. "
+                "Please either load a VPE file or generate VPEs from reference images."
             )
             return []
         
@@ -1260,7 +1311,7 @@ class DeployGeneratorDialog(QDialog):
                                             retina_masks=self.task == "segment")
 
         return [results]
-        
+
     def _apply_model(self, inputs):
         """
         Apply the model to the target inputs. This method handles both image-based 
@@ -1278,21 +1329,9 @@ class DeployGeneratorDialog(QDialog):
         
         # Check if the user is using VPE or Reference Images
         if self.reference_method_combo_box.currentText() == "VPE":
-            # Check if we have any VPEs available (imported or reference-generated)
-            has_vpes = bool(self.imported_vpes or self.reference_vpes)
-            
-            # If we have reference images selected but no imported VPEs yet,
-            # warn the user only if we also don't have any reference images
-            if not has_vpes and not references_dict:
-                QMessageBox.warning(
-                    self,
-                    "No VPEs Available",
-                    "No VPEs available for prediction. Please either load a VPE file or select reference images."
-                )
-                return []
-                
-            # Use the VPE method, which will combine imported and reference VPEs
-            results = self._apply_model_using_vpe(inputs, references_dict)
+            # The VPE method will use pre-loaded/pre-generated VPEs.
+            # The internal checks for whether any VPEs exist are now inside _apply_model_using_vpe.
+            results = self._apply_model_using_vpe(inputs)
         else:  
             # Use Reference Images method - requires reference images
             if not references_dict:
@@ -1405,16 +1444,9 @@ class DeployGeneratorDialog(QDialog):
         
     def show_vpe(self):
         """
-        Show a visualization of the VPEs using PyQtGraph.
-        This method now always recalculates VPEs from the currently highlighted reference images.
+        Show a visualization of the currently stored VPEs using PyQtGraph.
         """
-        # Set cursor to busy while loading VPEs
-        QApplication.setOverrideCursor(Qt.WaitCursor)
-        
         try:
-            # Always sync with the live UI selection before visualizing.
-            self.update_stashed_references_from_ui()
-
             vpes_with_source = []
 
             # 1. Add any VPEs that were loaded from a file
@@ -1422,37 +1454,35 @@ class DeployGeneratorDialog(QDialog):
                 for vpe in self.imported_vpes:
                     vpes_with_source.append((vpe, "Import"))
 
-            # 2. Get the currently selected reference images from the stashed list
-            references_dict = self._get_references()
-
-            # 3. If there are reference images, calculate their VPEs and add with source type
-            if references_dict:
-                self.reload_model()
-                new_reference_vpes = self.references_to_vpe(references_dict, update_reference_vpes=True)
-                if new_reference_vpes:
-                    for vpe in new_reference_vpes:
-                        vpes_with_source.append((vpe, "Reference"))
+            # 2. Add any pre-generated VPEs from reference images
+            if self.reference_vpes:
+                for vpe in self.reference_vpes:
+                    vpes_with_source.append((vpe, "Reference"))
 
-            # 4. Check if there is anything to visualize
+            # 3. Check if there is anything to visualize
             if not vpes_with_source:
                 QMessageBox.warning(
                     self,
                     "No VPEs Available",
-                    "No VPEs available to visualize. Please either load a VPE file or select reference images."
+                    "No VPEs available to visualize. Please load a VPE file or generate VPEs from references first."
                 )
                 return
 
-            # 5. Create the visualization dialog, passing the list of tuples
+            # 4. Create the visualization dialog
             all_vpe_tensors = [vpe for vpe, source in vpes_with_source]
             averaged_vpe = torch.cat(all_vpe_tensors).mean(dim=0, keepdim=True)
             final_vpe = torch.nn.functional.normalize(averaged_vpe, p=2, dim=-1)
 
+            QApplication.setOverrideCursor(Qt.WaitCursor)
+            
             dialog = VPEVisualizationDialog(vpes_with_source, final_vpe, self)
+            QApplication.restoreOverrideCursor()
+            
             dialog.exec_()
             
-        finally:
-            # Always restore cursor, even if an exception occurs
+        except Exception as e:
             QApplication.restoreOverrideCursor()
+            QMessageBox.critical(self, "Error Visualizing VPE", f"An error occurred: {str(e)}")
         
     def deactivate_model(self):
         """

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,
@@ -171,7 +169,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)
 
@@ -322,6 +320,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
         """
@@ -334,9 +333,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.
@@ -368,11 +406,13 @@ 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,
             )
-
+            # Finish the progress bar
+            progress_bar.finish_progress()
+            # Update the status bar
             self.status_bar.setText(f"Loaded ({self.model_path}")
             QMessageBox.information(self.annotation_window, "Model Loaded", "Model loaded successfully")
 
@@ -388,7 +428,7 @@ class DeployPredictorDialog(QDialog):
             progress_bar.stop_progress()
             progress_bar.close()
             progress_bar = None
-
+            
     def resize_image(self, image):
         """
         Resize the image to the specified size.
@@ -503,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(),
@@ -573,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:
 
@@ -581,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()

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,