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

coralnet_toolbox/SeeAnything/QtDeployPredictor.py

@@ -2,23 +2,23 @@ import warnings
 
 import os
 import gc
-import ujson as json
 
 import numpy as np
 
-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,
-                             QWidget, QLineEdit, QFileDialog)
+import torch
+from torch.cuda import empty_cache
+from ultralytics.utils import ops
 
 from ultralytics import YOLOE
 from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
 from ultralytics.models.yolo.yoloe import YOLOEVPDetectPredictor
 
-from torch.cuda import empty_cache
-from ultralytics.utils import ops
+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,
+                             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,83 +130,18 @@ 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)
 
-    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 setup_parameters_layout(self):
         """
         Setup parameter control section in a group box.
@@ -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.
@@ -517,9 +434,6 @@ class DeployPredictorDialog(QDialog):
             # Open the image using rasterio
             image = rasterio_to_numpy(self.main_window.image_window.rasterio_images[image_path])
 
-        # Preprocess the image
-        # image = preprocess_image(image)
-
         # Save the original image
         self.original_image = image
         self.image_path = image_path
@@ -529,19 +443,57 @@ class DeployPredictorDialog(QDialog):
             self.resized_image = self.resize_image(image)
         else:
             self.resized_image = image
+            
+    def scale_prompts(self, bboxes, masks=None):
+        """
+        Scale the bounding boxes and masks to the resized image.
+        """
+        # 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]
+        bboxes[:, 1] = (bboxes[:, 1] / self.original_image.shape[0]) * self.resized_image.shape[0]
+        bboxes[:, 2] = (bboxes[:, 2] / self.original_image.shape[1]) * self.resized_image.shape[1]
+        bboxes[:, 3] = (bboxes[:, 3] / self.original_image.shape[0]) * self.resized_image.shape[0]
+
+        # Set the predictor
+        self.task = self.task_dropdown.currentText()
 
-    def predict_from_prompts(self, bboxes):
+        # Create a visual dictionary
+        visual_prompts = {
+            'bboxes': np.array(bboxes),
+            'cls': np.zeros(len(bboxes))
+        }
+        if self.task == 'segment':
+            if masks:
+                scaled_masks = []
+                for mask in masks:
+                    scaled_mask = np.array(mask, dtype=np.float32)
+                    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)
+                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]]))
+                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:
-            bbox (np.ndarray): The bounding boxes to use as prompts.
+            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,
+            QMessageBox.critical(self.annotation_window, 
                                  "Model Not Loaded",
                                  "Model not loaded, cannot make predictions")
             return None
@@ -549,28 +501,14 @@ class DeployPredictorDialog(QDialog):
         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]
-        bboxes[:, 1] = (bboxes[:, 1] / self.original_image.shape[0]) * self.resized_image.shape[0]
-        bboxes[:, 2] = (bboxes[:, 2] / self.original_image.shape[1]) * self.resized_image.shape[1]
-        bboxes[:, 3] = (bboxes[:, 3] / self.original_image.shape[0]) * self.resized_image.shape[0]
-
-        # Create a visual dictionary
-        visuals = {
-            'bboxes': np.array(bboxes),
-            'cls': np.zeros(len(bboxes))  # TODO figure this out
-        }
-
-        # Set the predictor
-        self.task = self.task_dropdown.currentText()
-        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(),
@@ -590,7 +528,7 @@ class DeployPredictorDialog(QDialog):
 
         return results
 
-    def predict_from_annotations(self, refer_image, refer_label, refer_annotations, target_images):
+    def predict_from_annotations(self, refer_image, refer_label, refer_bboxes, refer_masks, target_images):
         """"""
         # Create a class mapping
         class_mapping = {0: refer_label}
@@ -605,15 +543,30 @@ class DeployPredictorDialog(QDialog):
             max_area_thresh=self.main_window.get_area_thresh_max()
         )
 
-        # Create a visual dictionary
-        visuals = {
-            'bboxes': np.array(refer_annotations),
-            'cls': np.zeros(len(refer_annotations))
-        }
-
-        # Set the predictor
-        self.task = self.task_dropdown.currentText()
-        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)
@@ -627,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/SeeAnything/QtTrainModel.py

@@ -13,7 +13,7 @@ from ultralytics.models.yolo.yoloe import YOLOEPESegTrainer
 from PyQt5.QtCore import Qt, QThread, pyqtSignal
 from PyQt5.QtWidgets import (QFileDialog, QScrollArea, QMessageBox, QCheckBox, QWidget, QVBoxLayout,
                              QLabel, QLineEdit, QDialog, QHBoxLayout, QPushButton, QComboBox, QSpinBox,
-                             QFormLayout, QTabWidget, QDoubleSpinBox, QGroupBox, )
+                             QFormLayout, QTabWidget, QDoubleSpinBox, QGroupBox)
 
 from torch.cuda import empty_cache
 
@@ -109,7 +109,7 @@ class TrainModelWorker(QThread):
 
             # Train the model
             self.model.train(**self.params,
-                             trainer=YOLOEPESegTrainer,
+                             trainer=YOLOEVPTrainer  , #YOLOEPESegTrainer,
                              device=self.device)
 
             # Post-run cleanup
@@ -222,6 +222,13 @@ class TrainModelDialog(QDialog):
         # Task specific parameters
         self.imgsz = 640
         self.batch = 4
+        
+        # Training parameters with defaults for linear-probing
+        self._lr0 = 1e-3
+        self._warmup_bias_lr = 0.0
+        self._weight_decay = 0.025
+        self._momentum = 0.9
+        self._close_mosaic = 0  # Default for linear-probing
 
         # Create the layout
         self.layout = QVBoxLayout(self)
@@ -232,6 +239,8 @@ class TrainModelDialog(QDialog):
         self.setup_dataset_layout()
         # Create the model layout (new)
         self.setup_model_layout()
+        # Create the output parameters layout
+        self.setup_output_parameters_layout()
         # Create and set up the parameters layout
         self.setup_parameters_layout()
         # Create the buttons layout
@@ -327,6 +336,29 @@ class TrainModelDialog(QDialog):
         layout.addWidget(tab_widget)
         group_box.setLayout(layout)
         self.layout.addWidget(group_box)
+        
+    def setup_output_parameters_layout(self):
+        """
+        Set up the layout and widgets for output parameters (project directory and name).
+        """
+        group_box = QGroupBox("Output")
+        layout = QFormLayout()
+
+        # Project
+        self.project_edit = QLineEdit()
+        self.project_button = QPushButton("Browse...")
+        self.project_button.clicked.connect(self.browse_project_dir)
+        project_layout = QHBoxLayout()
+        project_layout.addWidget(self.project_edit)
+        project_layout.addWidget(self.project_button)
+        layout.addRow("Project:", project_layout)
+
+        # Name
+        self.name_edit = QLineEdit()
+        layout.addRow("Name:", self.name_edit)
+
+        group_box.setLayout(layout)
+        self.layout.addWidget(group_box)
 
     def setup_parameters_layout(self):
         """
@@ -352,19 +384,6 @@ class TrainModelDialog(QDialog):
         group_layout = QVBoxLayout(group_box)
         group_layout.addWidget(scroll_area)
 
-        # Project
-        self.project_edit = QLineEdit()
-        self.project_button = QPushButton("Browse...")
-        self.project_button.clicked.connect(self.browse_project_dir)
-        project_layout = QHBoxLayout()
-        project_layout.addWidget(self.project_edit)
-        project_layout.addWidget(self.project_button)
-        form_layout.addRow("Project:", project_layout)
-
-        # Name
-        self.name_edit = QLineEdit()
-        form_layout.addRow("Name:", self.name_edit)
-
         # Fine-tune or Linear Probing
         self.training_mode = QComboBox()
         self.training_mode.addItems(["linear-probe", "fine-tune"])
@@ -381,11 +400,60 @@ class TrainModelDialog(QDialog):
 
         # Patience
         self.patience_spinbox = QSpinBox()
-        self.patience_spinbox.setMinimum(1)
+        self.patience_spinbox.setMinimum(0)  # Changed minimum to 0 to allow for 0 patience
         self.patience_spinbox.setMaximum(1000)
-        self.patience_spinbox.setValue(30)
+        self.patience_spinbox.setValue(0)  # Default for linear-probing
         form_layout.addRow("Patience:", self.patience_spinbox)
 
+        # Close Mosaic
+        self.close_mosaic_spinbox = QSpinBox()
+        self.close_mosaic_spinbox.setMinimum(0)
+        self.close_mosaic_spinbox.setMaximum(1000)
+        self.close_mosaic_spinbox.setValue(0)  # Default for linear-probing
+        form_layout.addRow("Close Mosaic:", self.close_mosaic_spinbox)
+        
+        # Optimizer
+        self.optimizer_combo = QComboBox()
+        self.optimizer_combo.addItems(["auto", "SGD", "Adam", "AdamW", "NAdam", "RAdam", "RMSProp"])
+        self.optimizer_combo.setCurrentText("AdamW")
+        form_layout.addRow("Optimizer:", self.optimizer_combo)
+
+        # Learning Rate (lr0)
+        self.lr0_spinbox = QDoubleSpinBox()
+        self.lr0_spinbox.setDecimals(6)
+        self.lr0_spinbox.setMinimum(0.000001)
+        self.lr0_spinbox.setMaximum(1.0)
+        self.lr0_spinbox.setSingleStep(0.0001)
+        self.lr0_spinbox.setValue(self._lr0)
+        form_layout.addRow("Learning Rate:", self.lr0_spinbox)
+
+        # Warmup Bias Learning Rate
+        self.warmup_bias_lr_spinbox = QDoubleSpinBox()
+        self.warmup_bias_lr_spinbox.setDecimals(6)
+        self.warmup_bias_lr_spinbox.setMinimum(0.0)
+        self.warmup_bias_lr_spinbox.setMaximum(1.0)
+        self.warmup_bias_lr_spinbox.setSingleStep(0.0001)
+        self.warmup_bias_lr_spinbox.setValue(self._warmup_bias_lr)
+        form_layout.addRow("Warmup Bias LR:", self.warmup_bias_lr_spinbox)
+
+        # Weight Decay
+        self.weight_decay_spinbox = QDoubleSpinBox()
+        self.weight_decay_spinbox.setDecimals(6)
+        self.weight_decay_spinbox.setMinimum(0.0)
+        self.weight_decay_spinbox.setMaximum(1.0)
+        self.weight_decay_spinbox.setSingleStep(0.001)
+        self.weight_decay_spinbox.setValue(self._weight_decay)
+        form_layout.addRow("Weight Decay:", self.weight_decay_spinbox)
+
+        # Momentum
+        self.momentum_spinbox = QDoubleSpinBox()
+        self.momentum_spinbox.setDecimals(2)
+        self.momentum_spinbox.setMinimum(0.0)
+        self.momentum_spinbox.setMaximum(1.0)
+        self.momentum_spinbox.setSingleStep(0.01)
+        self.momentum_spinbox.setValue(self._momentum)
+        form_layout.addRow("Momentum:", self.momentum_spinbox)
+
         # Imgsz
         self.imgsz_spinbox = QSpinBox()
         self.imgsz_spinbox.setMinimum(16)
@@ -393,10 +461,6 @@ class TrainModelDialog(QDialog):
         self.imgsz_spinbox.setValue(self.imgsz)
         form_layout.addRow("Image Size:", self.imgsz_spinbox)
 
-        # Multi Scale
-        self.multi_scale_combo = create_bool_combo()
-        form_layout.addRow("Multi-Scale:", self.multi_scale_combo)
-
         # Batch
         self.batch_spinbox = QSpinBox()
         self.batch_spinbox.setMinimum(1)
@@ -421,20 +485,6 @@ class TrainModelDialog(QDialog):
         self.save_period_spinbox.setMaximum(1000)
         self.save_period_spinbox.setValue(-1)
         form_layout.addRow("Save Period:", self.save_period_spinbox)
-
-        # Dropout
-        self.dropout_spinbox = QDoubleSpinBox()
-        self.dropout_spinbox.setMinimum(0.0)
-        self.dropout_spinbox.setMaximum(1.0)
-        self.dropout_spinbox.setValue(0.0)
-        form_layout.addRow("Dropout:", self.dropout_spinbox)
-
-        # Optimizer
-        self.optimizer_combo = QComboBox()
-        self.optimizer_combo.addItems(["auto", "SGD", "Adam", "AdamW", "NAdam", "RAdam", "RMSProp"])
-        self.optimizer_combo.setCurrentText("AdamW")
-        form_layout.addRow("Optimizer:", self.optimizer_combo)
-
         # Val
         self.val_combo = create_bool_combo()
         form_layout.addRow("Validation:", self.val_combo)
@@ -489,9 +539,32 @@ class TrainModelDialog(QDialog):
             # Fine-tune mode
             self.epochs_spinbox.setValue(80)
             self.patience_spinbox.setValue(20)
+            self.close_mosaic_spinbox.setValue(10)
+            self._close_mosaic = 10
+            
+            # These parameters stay the same for both modes
+            self.lr0_spinbox.setValue(1e-3)
+            self.warmup_bias_lr_spinbox.setValue(0.0)
+            self.weight_decay_spinbox.setValue(0.025)
+            self.momentum_spinbox.setValue(0.9)
+            
+            # Ensure optimizer is set to AdamW
+            self.optimizer_combo.setCurrentText("AdamW")
         else:
+            # Linear-probing mode
             self.epochs_spinbox.setValue(2)
-            self.patience_spinbox.setValue(1)
+            self.patience_spinbox.setValue(0)
+            self.close_mosaic_spinbox.setValue(0)
+            self._close_mosaic = 0
+            
+            # These parameters stay the same for both modes
+            self.lr0_spinbox.setValue(1e-3)
+            self.warmup_bias_lr_spinbox.setValue(0.0)
+            self.weight_decay_spinbox.setValue(0.025)
+            self.momentum_spinbox.setValue(0.9)
+            
+            # Ensure optimizer is set to AdamW
+            self.optimizer_combo.setCurrentText("AdamW")
 
     def load_model_combobox(self):
         """Load the model combobox with the available models."""
@@ -605,13 +678,16 @@ class TrainModelDialog(QDialog):
             'patience': self.patience_spinbox.value(),
             'batch': self.batch_spinbox.value(),
             'imgsz': self.imgsz_spinbox.value(),
-            'multi_scale': self.multi_scale_combo.currentText() == "True",
+            'optimizer': self.optimizer_combo.currentText(),
+            'lr0': self.lr0_spinbox.value(),
+            'warmup_bias_lr': self.warmup_bias_lr_spinbox.value(),
+            'weight_decay': self.weight_decay_spinbox.value(),
+            'momentum': self.momentum_spinbox.value(),
+            'close_mosaic': self.close_mosaic_spinbox.value(),
             'save': self.save_combo.currentText() == "True",
             'save_period': self.save_period_spinbox.value(),
             'workers': self.workers_spinbox.value(),
-            'optimizer': self.optimizer_combo.currentText(),
             'verbose': self.verbose_combo.currentText() == "True",
-            'dropout': self.dropout_spinbox.value(),
             'val': self.val_combo.currentText() == "True",
             'exist_ok': True,
             'plots': True,
@@ -644,12 +720,6 @@ class TrainModelDialog(QDialog):
                         except ValueError:
                             params[name] = value
 
-        params['lr0'] = 1e-3 if 'lr0' not in params else params['lr0']
-        params['warmup_bias_lr'] = 0.0 if 'warmup_bias_lr' not in params else params['warmup_bias_lr']
-        params['weight_decay'] = 0.025 if 'weight_decay' not in params else params['weight_decay']
-        params['momentum'] = 0.9 if 'momentum' not in params else params['momentum']
-        params['close_mosaic'] = 10 if 'close_mosaic' not in params else params['close_mosaic']
-
         # Return the dictionary of parameters
         return params
 

coralnet_toolbox/SeeAnything/__init__.py

@@ -2,9 +2,11 @@
 from .QtTrainModel import TrainModelDialog
 from .QtBatchInference import BatchInferenceDialog
 from .QtDeployPredictor import DeployPredictorDialog
+from .QtDeployGenerator import DeployGeneratorDialog
 
 __all__ = [
     'TrainModelDialog',
     'BatchInferenceDialog'
     'DeployPredictorDialog',
+    'DeployGeneratorDialog',
 ]