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

coralnet_toolbox/Explorer/QtExplorer.py

@@ -16,7 +16,7 @@ from PyQt5.QtWidgets import (QVBoxLayout, QHBoxLayout, QGraphicsView, QScrollAre
                              QGraphicsScene, QPushButton, QComboBox, QLabel, QWidget,
                              QMainWindow, QSplitter, QGroupBox, QSlider, QMessageBox,
                              QApplication, QGraphicsRectItem, QRubberBand, QMenu,
-                             QWidgetAction, QToolButton, QAction)
+                             QWidgetAction, QToolButton, QAction, QDoubleSpinBox)
 
 from coralnet_toolbox.Explorer.QtFeatureStore import FeatureStore
 from coralnet_toolbox.Explorer.QtDataItem import AnnotationDataItem
@@ -28,6 +28,7 @@ from coralnet_toolbox.Explorer.QtSettingsWidgets import UncertaintySettingsWidge
 from coralnet_toolbox.Explorer.QtSettingsWidgets import MislabelSettingsWidget
 from coralnet_toolbox.Explorer.QtSettingsWidgets import EmbeddingSettingsWidget
 from coralnet_toolbox.Explorer.QtSettingsWidgets import AnnotationSettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import DuplicateSettingsWidget
 
 from coralnet_toolbox.Annotations.QtRectangleAnnotation import RectangleAnnotation
 
@@ -68,6 +69,8 @@ class EmbeddingViewer(QWidget):
     mislabel_parameters_changed = pyqtSignal(dict) 
     find_uncertain_requested = pyqtSignal()
     uncertainty_parameters_changed = pyqtSignal(dict)
+    find_duplicates_requested = pyqtSignal()
+    duplicate_parameters_changed = pyqtSignal(dict)
 
     def __init__(self, parent=None):
         """Initialize the EmbeddingViewer widget."""
@@ -193,8 +196,37 @@ class EmbeddingViewer(QWidget):
         
         uncertainty_settings_widget.parameters_changed.connect(self.uncertainty_parameters_changed.emit)
         toolbar_layout.addWidget(self.find_uncertain_button)
+        
+        # Create a QToolButton for duplicate detection
+        self.find_duplicates_button = QToolButton()
+        self.find_duplicates_button.setText("Find Duplicates")
+        self.find_duplicates_button.setToolTip(
+            "Find annotations that are likely duplicates based on feature similarity."
+        )
+        self.find_duplicates_button.setPopupMode(QToolButton.MenuButtonPopup)
+        self.find_duplicates_button.setToolButtonStyle(Qt.ToolButtonTextOnly)
+        self.find_duplicates_button.setStyleSheet(
+            "QToolButton::menu-indicator { "
+            "subcontrol-position: right center; "
+            "subcontrol-origin: padding; "
+            "left: -4px; }"
+        )
+
+        run_duplicates_action = QAction("Find Duplicates", self)
+        run_duplicates_action.triggered.connect(self.find_duplicates_requested.emit)
+        self.find_duplicates_button.setDefaultAction(run_duplicates_action)
+
+        duplicate_settings_widget = DuplicateSettingsWidget()
+        duplicate_menu = QMenu(self)
+        duplicate_widget_action = QWidgetAction(duplicate_menu)
+        duplicate_widget_action.setDefaultWidget(duplicate_settings_widget)
+        duplicate_menu.addAction(duplicate_widget_action)
+        self.find_duplicates_button.setMenu(duplicate_menu)
+        
+        duplicate_settings_widget.parameters_changed.connect(self.duplicate_parameters_changed.emit)
+        toolbar_layout.addWidget(self.find_duplicates_button)
     
-        # Add a strech and separator
+        # Add a stretch and separator
         toolbar_layout.addStretch()
         toolbar_layout.addWidget(self._create_separator())
 
@@ -293,6 +325,7 @@ class EmbeddingViewer(QWidget):
 
         self.find_mislabels_button.setEnabled(points_exist)
         self.find_uncertain_button.setEnabled(points_exist and self.is_uncertainty_analysis_available)
+        self.find_duplicates_button.setEnabled(points_exist)
         self.center_on_selection_button.setEnabled(points_exist and selection_exists)
 
         if self.isolated_mode:
@@ -348,6 +381,7 @@ class EmbeddingViewer(QWidget):
         self.center_on_selection_button.setEnabled(False)  # Disable center button
         self.find_mislabels_button.setEnabled(False)
         self.find_uncertain_button.setEnabled(False)
+        self.find_duplicates_button.setEnabled(False)
 
         self.isolate_button.show()
         self.isolate_button.setEnabled(False)
@@ -851,7 +885,7 @@ class AnnotationViewer(QWidget):
                     
                 # Show resize handles for Rectangle annotations
                 if isinstance(annotation_to_select, RectangleAnnotation):
-                    explorer.annotation_window.set_selected_tool('select') # Accidently unselects in AnnotationWindow
+                    explorer.annotation_window.set_selected_tool('select')  # Accidentally unselects in AnnotationWindow
                     explorer.annotation_window.select_annotation(annotation_to_select, quiet_mode=True)
                     select_tool = explorer.annotation_window.tools.get('select')
 
@@ -1584,6 +1618,7 @@ class ExplorerWindow(QMainWindow):
         self.mislabel_params = {'k': 20, 'threshold': 0.6}
         self.uncertainty_params = {'confidence': 0.6, 'margin': 0.1}
         self.similarity_params = {'k': 30}
+        self.duplicate_params = {'threshold': 0.05}
         
         self.data_item_cache = {}  # Cache for AnnotationDataItem objects
 
@@ -1670,37 +1705,27 @@ class ExplorerWindow(QMainWindow):
             if child.widget():
                 child.widget().setParent(None)
 
-        # Lazily initialize the settings and viewer widgets if they haven't been created yet.
-        # This ensures that the widgets are only created once per ExplorerWindow instance.
-
-        # Annotation settings panel (filters by image, type, label)
+        # Lazily initialize the settings and viewer widgets
         if self.annotation_settings_widget is None:
             self.annotation_settings_widget = AnnotationSettingsWidget(self.main_window, self)
-
-        # Model selection panel (choose feature extraction model)
         if self.model_settings_widget is None:
             self.model_settings_widget = ModelSettingsWidget(self.main_window, self)
-
-        # Embedding settings panel (choose dimensionality reduction method)
         if self.embedding_settings_widget is None:
             self.embedding_settings_widget = EmbeddingSettingsWidget(self.main_window, self)
-
-        # Annotation viewer (shows annotation image crops in a grid)
         if self.annotation_viewer is None:
             self.annotation_viewer = AnnotationViewer(self)
-
-        # Embedding viewer (shows 2D embedding scatter plot)
         if self.embedding_viewer is None:
             self.embedding_viewer = EmbeddingViewer(self)
 
+        # Horizontal layout for the three settings panels (original horizontal layout)
         top_layout = QHBoxLayout()
         top_layout.addWidget(self.annotation_settings_widget, 2)
         top_layout.addWidget(self.model_settings_widget, 1)
         top_layout.addWidget(self.embedding_settings_widget, 1)
         top_container = QWidget()
         top_container.setLayout(top_layout)
-        self.main_layout.addWidget(top_container)
 
+        # Horizontal splitter for the two main viewer panels
         middle_splitter = QSplitter(Qt.Horizontal)
         annotation_group = QGroupBox("Annotation Viewer")
         annotation_layout = QVBoxLayout(annotation_group)
@@ -1712,7 +1737,19 @@ class ExplorerWindow(QMainWindow):
         embedding_layout.addWidget(self.embedding_viewer)
         middle_splitter.addWidget(embedding_group)
         middle_splitter.setSizes([500, 500])
-        self.main_layout.addWidget(middle_splitter, 1)
+
+        # Create a VERTICAL splitter to manage the height between the settings and viewers.
+        # This makes the top settings panel vertically resizable.
+        main_splitter = QSplitter(Qt.Vertical)
+        main_splitter.addWidget(top_container)
+        main_splitter.addWidget(middle_splitter)
+        
+        # Set initial heights to give the settings panel a bit more space by default
+        main_splitter.setSizes([250, 750]) 
+
+        # Add the new main splitter to the layout instead of the individual components
+        self.main_layout.addWidget(main_splitter, 1)
+
         self.main_layout.addWidget(self.label_window)
 
         self.buttons_layout = QHBoxLayout()
@@ -1744,6 +1781,8 @@ class ExplorerWindow(QMainWindow):
         self.model_settings_widget.selection_changed.connect(self.on_model_selection_changed)
         self.embedding_viewer.find_uncertain_requested.connect(self.find_uncertain_annotations)
         self.embedding_viewer.uncertainty_parameters_changed.connect(self.on_uncertainty_params_changed)
+        self.embedding_viewer.find_duplicates_requested.connect(self.find_duplicate_annotations)
+        self.embedding_viewer.duplicate_parameters_changed.connect(self.on_duplicate_params_changed)
         self.annotation_viewer.find_similar_requested.connect(self.find_similar_annotations)
         self.annotation_viewer.similarity_settings_widget.parameters_changed.connect(self.on_similarity_params_changed)
         
@@ -1887,6 +1926,12 @@ class ExplorerWindow(QMainWindow):
         """Updates the stored parameters for uncertainty analysis."""
         self.uncertainty_params = params
         print(f"Uncertainty parameters updated: {self.uncertainty_params}")
+
+    @pyqtSlot(dict)
+    def on_duplicate_params_changed(self, params):
+        """Updates the stored parameters for duplicate detection."""
+        self.duplicate_params = params
+        print(f"Duplicate detection parameters updated: {self.duplicate_params}")
         
     @pyqtSlot(dict)
     def on_similarity_params_changed(self, params):
@@ -2067,6 +2112,98 @@ class ExplorerWindow(QMainWindow):
 
         finally:
             QApplication.restoreOverrideCursor()
+
+    def find_duplicate_annotations(self):
+        """
+        Identifies annotations that are likely duplicates based on feature similarity.
+        It uses a nearest-neighbor approach in the high-dimensional feature space.
+        For each group of duplicates found, it selects all but one "original".
+        """
+        threshold = self.duplicate_params.get('threshold', 0.05)
+
+        if not self.embedding_viewer.points_by_id or len(self.embedding_viewer.points_by_id) < 2:
+            QMessageBox.information(self, 
+                                    "Not Enough Data", 
+                                    "This feature requires at least 2 points in the embedding viewer.")
+            return
+
+        items_in_view = list(self.embedding_viewer.points_by_id.values())
+        data_items_in_view = [p.data_item for p in items_in_view]
+
+        model_info = self.model_settings_widget.get_selected_model()
+        model_name, feature_mode = model_info if isinstance(model_info, tuple) else (model_info, "default")
+        sanitized_model_name = os.path.basename(model_name).replace(' ', '_')
+        sanitized_feature_mode = feature_mode.replace(' ', '_').replace('/', '_')
+        model_key = f"{sanitized_model_name}_{sanitized_feature_mode}"
+
+        # Make cursor busy
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        try:
+            index = self.feature_store._get_or_load_index(model_key)
+            if index is None:
+                QMessageBox.warning(self, "Error", "Could not find a valid feature index for the current model.")
+                return
+
+            features_dict, _ = self.feature_store.get_features(data_items_in_view, model_key)
+            if not features_dict:
+                QMessageBox.warning(self, "Error", "Could not retrieve features for the items in view.")
+                return
+
+            query_ann_ids = list(features_dict.keys())
+            query_vectors = np.array([features_dict[ann_id] for ann_id in query_ann_ids]).astype('float32')
+
+            # Find the 2 nearest neighbors for each vector. D = squared L2 distances.
+            D, I = index.search(query_vectors, 2)
+
+            # Use a Disjoint Set Union (DSU) data structure to group duplicates.
+            parent = {ann_id: ann_id for ann_id in query_ann_ids}
+            
+            # Helper functions for DSU
+            def find_set(v):
+                if v == parent[v]:
+                    return v
+                parent[v] = find_set(parent[v])
+                return parent[v]
+            
+            def unite_sets(a, b):
+                a = find_set(a)
+                b = find_set(b)
+                if a != b:
+                    parent[b] = a
+            
+            id_map = self.feature_store.get_faiss_index_to_annotation_id_map(model_key)
+
+            for i, ann_id in enumerate(query_ann_ids):
+                neighbor_faiss_idx = I[i, 1]  # The second result is the nearest neighbor
+                distance = D[i, 1]
+
+                if distance < threshold:
+                    neighbor_ann_id = id_map.get(neighbor_faiss_idx)
+                    if neighbor_ann_id and neighbor_ann_id in parent:
+                        unite_sets(ann_id, neighbor_ann_id)
+            
+            # Group annotations by their set representative
+            groups = {}
+            for ann_id in query_ann_ids:
+                root = find_set(ann_id)
+                if root not in groups:
+                    groups[root] = []
+                groups[root].append(ann_id)
+
+            copies_to_select = set()
+            for root_id, group_ids in groups.items():
+                if len(group_ids) > 1:
+                    # Sort IDs to consistently pick the same "original".
+                    # Sorting strings is reliable.
+                    sorted_ids = sorted(group_ids)
+                    # The first ID is the original, add the rest to the selection.
+                    copies_to_select.update(sorted_ids[1:])
+            
+            print(f"Found {len(copies_to_select)} duplicate annotations.")
+            self.embedding_viewer.render_selection_from_ids(copies_to_select)
+
+        finally:
+            QApplication.restoreOverrideCursor()
             
     def find_uncertain_annotations(self):
         """
@@ -2658,6 +2795,7 @@ class ExplorerWindow(QMainWindow):
             norm_y = (embedded_features[i, 1] - min_vals[1]) / range_vals[1] if range_vals[1] > 0 else 0.5
             item.embedding_x = (norm_x * scale_factor) - (scale_factor / 2)
             item.embedding_y = (norm_y * scale_factor) - (scale_factor / 2)
+            item.embedding_id = i
 
     def run_embedding_pipeline(self):
         """
@@ -2813,6 +2951,10 @@ class ExplorerWindow(QMainWindow):
             self.current_data_items = [
                 item for item in self.current_data_items if item.annotation.id not in deleted_ann_ids
             ]
+            # Also update the annotation viewer's list to keep it in sync
+            self.annotation_viewer.all_data_items = [
+                item for item in self.annotation_viewer.all_data_items if item.annotation.id not in deleted_ann_ids
+            ]
             for ann_id in deleted_ann_ids:
                 if ann_id in self.data_item_cache:
                     del self.data_item_cache[ann_id]

coralnet_toolbox/Explorer/QtSettingsWidgets.py

@@ -4,7 +4,7 @@ import warnings
 from PyQt5.QtCore import Qt, pyqtSignal, pyqtSlot
 from PyQt5.QtWidgets import (QVBoxLayout, QHBoxLayout, QPushButton, QComboBox, QLabel, 
                              QWidget, QGroupBox, QSlider, QListWidget, QTabWidget, 
-                             QLineEdit, QFileDialog, QFormLayout, QSpinBox)
+                             QLineEdit, QFileDialog, QFormLayout, QSpinBox, QDoubleSpinBox)
 
 from coralnet_toolbox.MachineLearning.Community.cfg import get_available_configs
 
@@ -189,6 +189,48 @@ class SimilaritySettingsWidget(QWidget):
             'k': self.k_spinbox.value()
         }
         
+        
+class DuplicateSettingsWidget(QWidget):
+    """Widget for configuring duplicate detection parameters."""
+    parameters_changed = pyqtSignal(dict)
+
+    def __init__(self, parent=None):
+        super(DuplicateSettingsWidget, self).__init__(parent)
+        layout = QVBoxLayout(self)
+        layout.setContentsMargins(10, 10, 10, 10)
+        
+        # Using a DoubleSpinBox for the distance threshold
+        self.threshold_spinbox = QDoubleSpinBox()
+        self.threshold_spinbox.setDecimals(3)
+        self.threshold_spinbox.setRange(0.0, 10.0)
+        self.threshold_spinbox.setSingleStep(0.01)
+        self.threshold_spinbox.setValue(0.1)  # Default value for squared L2 distance
+        self.threshold_spinbox.setToolTip(
+            "Similarity Threshold (Squared L2 Distance).\n"
+            "Lower values mean more similar.\n"
+            "A value of 0 means identical features."
+        )
+        
+        self.threshold_spinbox.valueChanged.connect(self._emit_parameters)
+
+        form_layout = QHBoxLayout()
+        form_layout.addWidget(QLabel("Threshold:"))
+        form_layout.addWidget(self.threshold_spinbox)
+        layout.addLayout(form_layout)
+
+    def _emit_parameters(self):
+        """Emits the current parameters."""
+        params = {
+            'threshold': self.threshold_spinbox.value()
+        }
+        self.parameters_changed.emit(params)
+
+    def get_parameters(self):
+        """Returns the current parameters as a dictionary."""
+        return {
+            'threshold': self.threshold_spinbox.value()
+        }
+
 
 class AnnotationSettingsWidget(QGroupBox):
     """Widget for filtering annotations by image, type, and label in a multi-column layout."""
@@ -200,7 +242,7 @@ class AnnotationSettingsWidget(QGroupBox):
         self.setup_ui()
 
     def setup_ui(self):
-        # The main layout is vertical, to hold the top columns, the stretch, and the bottom buttons
+        # The main layout is vertical, to hold the top columns and the bottom buttons
         layout = QVBoxLayout(self)
 
         # A horizontal layout to contain the filter columns
@@ -213,8 +255,8 @@ class AnnotationSettingsWidget(QGroupBox):
         images_column.addWidget(images_label)
 
         self.images_list = QListWidget()
-        self.images_list.setSelectionMode(QListWidget.MultiSelection)
-        self.images_list.setMaximumHeight(50)
+        self.images_list.setSelectionMode(QListWidget.ExtendedSelection)
+        self.images_list.setMinimumHeight(100)
 
         if hasattr(self.main_window, 'image_window') and hasattr(self.main_window.image_window, 'raster_manager'):
             for path in self.main_window.image_window.raster_manager.image_paths:
@@ -241,8 +283,8 @@ class AnnotationSettingsWidget(QGroupBox):
         type_column.addWidget(type_label)
 
         self.annotation_type_list = QListWidget()
-        self.annotation_type_list.setSelectionMode(QListWidget.MultiSelection)
-        self.annotation_type_list.setMaximumHeight(50)
+        self.annotation_type_list.setSelectionMode(QListWidget.ExtendedSelection)
+        self.annotation_type_list.setMinimumHeight(100)
         self.annotation_type_list.addItems(["PatchAnnotation",
                                             "RectangleAnnotation",
                                             "PolygonAnnotation",
@@ -269,8 +311,8 @@ class AnnotationSettingsWidget(QGroupBox):
         label_column.addWidget(label_label)
 
         self.label_list = QListWidget()
-        self.label_list.setSelectionMode(QListWidget.MultiSelection)
-        self.label_list.setMaximumHeight(50)
+        self.label_list.setSelectionMode(QListWidget.ExtendedSelection)
+        self.label_list.setMinimumHeight(100)
 
         if hasattr(self.main_window, 'label_window') and hasattr(self.main_window.label_window, 'labels'):
             for label in self.main_window.label_window.labels:
@@ -293,9 +335,6 @@ class AnnotationSettingsWidget(QGroupBox):
         # Add the horizontal layout of columns to the main vertical layout
         layout.addLayout(conditions_layout)
 
-        # Add a stretch item to push the columns to the top
-        layout.addStretch(1)
-
         # Bottom buttons layout with Apply and Clear buttons on the right
         bottom_layout = QHBoxLayout()
         bottom_layout.addStretch()  # Push buttons to the right
@@ -308,9 +347,12 @@ class AnnotationSettingsWidget(QGroupBox):
         self.clear_button.clicked.connect(self.clear_all_conditions)
         bottom_layout.addWidget(self.clear_button)
 
-        # Add the bottom buttons layout to the main layout, keeping it at the bottom
+        # Add the bottom buttons layout to the main layout
         layout.addLayout(bottom_layout)
 
+        # Add a stretch item to push all content to the top
+        layout.addStretch(1)
+
         # Set defaults
         self.set_defaults()
 
@@ -393,68 +435,66 @@ class AnnotationSettingsWidget(QGroupBox):
 
 
 class ModelSettingsWidget(QGroupBox):
-    """Widget containing model selection with tabs for different model sources."""
+    """Widget containing a structured, hierarchical model selection system."""
     selection_changed = pyqtSignal()
 
     def __init__(self, main_window, parent=None):
         super(ModelSettingsWidget, self).__init__("Model Settings", parent)
         self.main_window = main_window
         self.explorer_window = parent
+
+        # --- Data for hierarchical selection ---
+        self.standard_models_map = {
+            'YOLOv8': {'Nano': 'n', 'Small': 's', 'Medium': 'm', 'Large': 'l', 'X-Large': 'x'},
+            'YOLOv11': {'Nano': 'n', 'Small': 's', 'Medium': 'm', 'Large': 'l', 'X-Large': 'x'},
+            'YOLOv12': {'Nano': 'n', 'Small': 's', 'Medium': 'm', 'Large': 'l', 'X-Large': 'x'}
+        }
+        self.community_configs = get_available_configs(task='classify')
+
         self.setup_ui()
 
     def setup_ui(self):
         """Set up the UI with a tabbed interface for model selection."""
         main_layout = QVBoxLayout(self)
-
-        # --- Tabbed Interface for Model Selection ---
         self.tabs = QTabWidget()
 
-        # Tab 1: Select Model
+        # === Tab 1: Select Pre-defined Model ===
         model_select_tab = QWidget()
-        model_select_layout = QFormLayout(model_select_tab)
-        model_select_layout.setContentsMargins(5, 10, 5, 5)  # Add some top margin
-
-        self.model_combo = QComboBox()
-        self.model_combo.addItems(["Color Features"])
-        self.model_combo.insertSeparator(1)  # Add a separator
-
-        standard_models = [
-            'yolov8n-cls.pt',
-            'yolov8s-cls.pt',
-            'yolov8m-cls.pt',
-            'yolov8l-cls.pt',
-            'yolov8x-cls.pt',
-            'yolo11n-cls.pt',
-            'yolo11s-cls.pt',
-            'yolo11m-cls.pt',
-            'yolo11l-cls.pt',
-            'yolo11x-cls.pt',
-            'yolo12n-cls.pt',
-            'yolo12s-cls.pt',
-            'yolo12m-cls.pt',
-            'yolo12l-cls.pt',
-            'yolo12x-cls.pt'
-        ]
-
-        self.model_combo.addItems(standard_models)
-
-        community_configs = get_available_configs(task='classify')
-        if community_configs:
-            self.model_combo.insertSeparator(len(standard_models) + 2)
-            self.model_combo.addItems(list(community_configs.keys()))
-
-        self.model_combo.setCurrentText('Color Features')
-        model_select_layout.addRow("Model:", self.model_combo)
+        self.model_select_layout = QFormLayout(model_select_tab)
+        self.model_select_layout.setContentsMargins(5, 10, 5, 5)
+
+        # 1. Main Category ComboBox
+        self.category_combo = QComboBox()
+        self.category_combo.addItems(["Color Features", "Standard Model", "Community Model"])
+        self.model_select_layout.addRow("Category:", self.category_combo)
+
+        # 2. Standard Model Options (initially hidden)
+        self.family_combo = QComboBox()
+        self.family_combo.addItems(self.standard_models_map.keys())
+        self.size_combo = QComboBox()
+        self.family_combo.currentTextChanged.connect(self._update_size_combo)
+        self._update_size_combo(self.family_combo.currentText())
+        
+        self.standard_model_widgets = [QLabel("Family:"), self.family_combo, QLabel("Size:"), self.size_combo]
+        self.model_select_layout.addRow(self.standard_model_widgets[0], self.standard_model_widgets[1])
+        self.model_select_layout.addRow(self.standard_model_widgets[2], self.standard_model_widgets[3])
+
+        # 3. Community Model Options (initially hidden)
+        self.community_combo = QComboBox()
+        if self.community_configs:
+            self.community_combo.addItems(list(self.community_configs.keys()))
+        self.community_model_widgets = [QLabel("Model:"), self.community_combo]
+        self.model_select_layout.addRow(self.community_model_widgets[0], self.community_model_widgets[1])
         
         self.tabs.addTab(model_select_tab, "Select Model")
 
-        # Tab 2: Existing Model from File
+        # === Tab 2: Existing Model from File ===
         model_existing_tab = QWidget()
         model_existing_layout = QFormLayout(model_existing_tab)
         model_existing_layout.setContentsMargins(5, 10, 5, 5)
 
         self.model_path_edit = QLineEdit()
-        self.model_path_edit.setPlaceholderText("Path to a existing .pt model file...")
+        self.model_path_edit.setPlaceholderText("Path to a compatible .pt model file...")
         browse_button = QPushButton("Browse...")
         browse_button.clicked.connect(self.browse_for_model)
 
@@ -463,82 +503,116 @@ class ModelSettingsWidget(QGroupBox):
         path_layout.addWidget(browse_button)
         model_existing_layout.addRow("Model Path:", path_layout)
         
+        help_label = QLabel("Note: Select an Ultralytics model (.pt).")
+        help_label.setStyleSheet("color: gray; font-style: italic;")
+        model_existing_layout.addRow("", help_label)
+
         self.tabs.addTab(model_existing_tab, "Use Existing Model")
 
         main_layout.addWidget(self.tabs)
-        
-        # Connect all relevant widgets to a single slot that emits the new signal
-        self.model_combo.currentTextChanged.connect(self._on_selection_changed)
-        self.tabs.currentChanged.connect(self._on_selection_changed)
-        self.model_path_edit.textChanged.connect(self._on_selection_changed)
-        
-        # Add feature extraction mode selection outside of tabs
+
+        # === Feature Extraction Mode (Reverted to bottom) ===
         feature_mode_layout = QFormLayout()
         self.feature_mode_combo = QComboBox()
         self.feature_mode_combo.addItems(["Predictions", "Embed Features"])
-        self.feature_mode_combo.currentTextChanged.connect(self._on_selection_changed)
         feature_mode_layout.addRow("Feature Mode:", self.feature_mode_combo)
         main_layout.addLayout(feature_mode_layout)
+
+        # --- Connect Signals ---
+        self.category_combo.currentTextChanged.connect(self._on_category_changed)
+        self.tabs.currentChanged.connect(self._on_selection_changed)
+        for widget in [self.category_combo, self.family_combo, self.size_combo, 
+                       self.community_combo, self.model_path_edit, self.feature_mode_combo]:
+            if isinstance(widget, QComboBox):
+                widget.currentTextChanged.connect(self._on_selection_changed)
+            elif isinstance(widget, QLineEdit):
+                widget.textChanged.connect(self._on_selection_changed)
+        
+        # --- Initial State ---
+        self._on_category_changed(self.category_combo.currentText())
+        self._on_selection_changed()
+
+    @pyqtSlot(str)
+    def _update_size_combo(self, family):
+        """Populate the size combo box based on the selected family."""
+        self.size_combo.clear()
+        if family in self.standard_models_map:
+            self.size_combo.addItems(self.standard_models_map[family].keys())
+
+    @pyqtSlot(str)
+    def _on_category_changed(self, category):
+        """Show or hide sub-option widgets based on the selected category."""
+        is_standard = (category == "Standard Model")
+        is_community = (category == "Community Model")
+        
+        for widget in self.standard_model_widgets:
+            widget.setVisible(is_standard)
+        for widget in self.community_model_widgets:
+            widget.setVisible(is_community)
         
-        # Initialize the feature mode state and emit the first signal
         self._on_selection_changed()
 
     def browse_for_model(self):
         """Open a file dialog to browse for model files."""
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(
-            self,
-            "Select Model File",
-            "",
-            "PyTorch Models (*.pt);;All Files (*)",
-            options=options
-        )
+            self, "Select Model File", "", "PyTorch Models (*.pt);;All Files (*)", options=options)
         if file_path:
             self.model_path_edit.setText(file_path)
             
     @pyqtSlot()
     def _on_selection_changed(self):
-        """Central slot to handle any change in model selection and emit a single signal."""
+        """Central slot to handle any change and emit a single signal."""
         self._update_feature_mode_state()
         self.selection_changed.emit()
 
-    def _update_feature_mode_state(self, *args):
-        """Update the enabled state of the feature mode field based on the current model selection."""
-        current_tab_index = self.tabs.currentIndex()
+    def _update_feature_mode_state(self):
+        """Update the enabled state and tooltip of the feature mode field."""
         is_color_features = False
+        current_tab_index = self.tabs.currentIndex()
         
         if current_tab_index == 0:
-            # Select Model tab - check if Color Features is selected
-            current_model = self.model_combo.currentText()
-            is_color_features = current_model == "Color Features"
-        elif current_tab_index == 1:
-            # Use Existing Model tab - feature mode should always be enabled
-            is_color_features = False
+            is_color_features = (self.category_combo.currentText() == "Color Features")
         
-        # Enable feature mode only if not Color Features
         self.feature_mode_combo.setEnabled(not is_color_features)
         
-        # Update the tooltip based on state
         if is_color_features:
-            self.feature_mode_combo.setToolTip("Feature Mode is not available for Color Features")
+            self.feature_mode_combo.setToolTip("Feature Mode is not applicable for Color Features.")
         else:
-            self.feature_mode_combo.setToolTip("Select the feature extraction mode")
+            self.feature_mode_combo.setToolTip(
+                "Choose 'Predictions' for class probabilities (for uncertainty analysis)\n"
+                "or 'Embed Features' for a general-purpose feature vector."
+            )
     
     def get_selected_model(self):
         """Get the currently selected model name/path and feature mode."""
         current_tab_index = self.tabs.currentIndex()
+        model_name = ""
         
-        # Get model name/path and feature mode based on the active tab
         if current_tab_index == 0:
-            model_name = self.model_combo.currentText()
+            category = self.category_combo.currentText()
+            if category == "Color Features":
+                model_name = "Color Features"
+            elif category == "Standard Model":
+                family_text = self.family_combo.currentText()
+                size_text = self.size_combo.currentText()
+                
+                # Add a guard clause to prevent crashing if a combo is empty.
+                if not family_text or not size_text:
+                    return "", "N/A"  # Return a safe default
+                
+                family_key = family_text.lower().replace('-', '')
+                size_key = self.standard_models_map[family_text][size_text]
+                model_name = f"{family_key}{size_key}-cls.pt"
+
+            elif category == "Community Model":
+                model_name = self.community_combo.currentText()
         elif current_tab_index == 1:
             model_name = self.model_path_edit.text()
-        else:
-            return "", None
         
         feature_mode = self.feature_mode_combo.currentText() if self.feature_mode_combo.isEnabled() else "N/A"
         return model_name, feature_mode
-
+    
 
 class EmbeddingSettingsWidget(QGroupBox):
     """Widget containing settings with tabs for models and embedding."""