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

coralnet_toolbox/Explorer/QtSettingsWidgets.py

@@ -7,6 +7,8 @@ from PyQt5.QtWidgets import (QVBoxLayout, QHBoxLayout, QPushButton, QComboBox, Q
                              QLineEdit, QFileDialog, QFormLayout, QSpinBox, QDoubleSpinBox)
 
 from coralnet_toolbox.MachineLearning.Community.cfg import get_available_configs
+from coralnet_toolbox.Explorer.transformer_models import TRANSFORMER_MODELS
+from coralnet_toolbox.Explorer.yolo_models import YOLO_MODELS
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -444,15 +446,54 @@ class ModelSettingsWidget(QGroupBox):
         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'}
-        }
+        # Convert YOLO_MODELS to a hierarchical structure for the UI
+        self.standard_models_map = self._create_hierarchical_model_map()
         self.community_configs = get_available_configs(task='classify')
+        
+        # --- Models for feature extraction ---
+        self.transformer_models = TRANSFORMER_MODELS
+        self.yolo_models = YOLO_MODELS
 
         self.setup_ui()
 
+    def _create_hierarchical_model_map(self):
+        """Convert the flat YOLO_MODELS dictionary to a hierarchical structure."""
+        hierarchical_map = {}
+        
+        # Process each model in YOLO_MODELS
+        for display_name, model_file in YOLO_MODELS.items():
+            # Parse the family and size from the display name
+            # Expected format: "YOLOvX (Size)"
+            parts = display_name.split(' ')
+            if len(parts) >= 2:
+                family = parts[0]  # "YOLOv8", "YOLOv11", etc.
+                size = parts[1].strip('()')  # "Nano", "Small", etc.
+                
+                # Create the family entry if it doesn't exist
+                if family not in hierarchical_map:
+                    hierarchical_map[family] = {}
+                
+                # Extract the size code from the model filename
+                # For example, get 'n' from 'yolov8n-cls.pt'
+                if model_file.startswith(family.lower()):
+                    size_code = model_file[len(family.lower())]
+                else:
+                    # Default fallback if parsing fails
+                    size_code = size[0].lower()
+                
+                # Add the size entry to the family
+                hierarchical_map[family][size] = size_code
+        
+        # If no models were processed, fallback to the default structure
+        if not hierarchical_map:
+            hierarchical_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'}
+            }
+            
+        return hierarchical_map
+
     def setup_ui(self):
         """Set up the UI with a tabbed interface for model selection."""
         main_layout = QVBoxLayout(self)
@@ -465,7 +506,7 @@ class ModelSettingsWidget(QGroupBox):
 
         # 1. Main Category ComboBox
         self.category_combo = QComboBox()
-        self.category_combo.addItems(["Color Features", "Standard Model", "Community Model"])
+        self.category_combo.addItems(["Color Features", "Standard Model", "Community Model", "Transformer Model"])
         self.model_select_layout.addRow("Category:", self.category_combo)
 
         # 2. Standard Model Options (initially hidden)
@@ -486,6 +527,12 @@ class ModelSettingsWidget(QGroupBox):
         self.community_model_widgets = [QLabel("Model:"), self.community_combo]
         self.model_select_layout.addRow(self.community_model_widgets[0], self.community_model_widgets[1])
         
+        # 4. Transformer Model Options (initially hidden)
+        self.transformer_combo = QComboBox()
+        self.transformer_combo.addItems(list(self.transformer_models.keys()))
+        self.transformer_model_widgets = [QLabel("Model:"), self.transformer_combo]
+        self.model_select_layout.addRow(self.transformer_model_widgets[0], self.transformer_model_widgets[1])
+        
         self.tabs.addTab(model_select_tab, "Select Model")
 
         # === Tab 2: Existing Model from File ===
@@ -515,6 +562,7 @@ class ModelSettingsWidget(QGroupBox):
         feature_mode_layout = QFormLayout()
         self.feature_mode_combo = QComboBox()
         self.feature_mode_combo.addItems(["Predictions", "Embed Features"])
+        self.feature_mode_combo.setCurrentText("Embed Features")  # Set default to Embed Features
         feature_mode_layout.addRow("Feature Mode:", self.feature_mode_combo)
         main_layout.addLayout(feature_mode_layout)
 
@@ -522,7 +570,7 @@ class ModelSettingsWidget(QGroupBox):
         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]:
+                       self.community_combo, self.transformer_combo, self.model_path_edit, self.feature_mode_combo]:
             if isinstance(widget, QComboBox):
                 widget.currentTextChanged.connect(self._on_selection_changed)
             elif isinstance(widget, QLineEdit):
@@ -544,11 +592,14 @@ class ModelSettingsWidget(QGroupBox):
         """Show or hide sub-option widgets based on the selected category."""
         is_standard = (category == "Standard Model")
         is_community = (category == "Community Model")
+        is_transformer = (category == "Transformer Model")
         
         for widget in self.standard_model_widgets:
             widget.setVisible(is_standard)
         for widget in self.community_model_widgets:
             widget.setVisible(is_community)
+        for widget in self.transformer_model_widgets:
+            widget.setVisible(is_transformer)
         
         self._on_selection_changed()
 
@@ -569,15 +620,29 @@ class ModelSettingsWidget(QGroupBox):
     def _update_feature_mode_state(self):
         """Update the enabled state and tooltip of the feature mode field."""
         is_color_features = False
+        is_transformer = False
+        is_community = False  # Add a new flag for community models
         current_tab_index = self.tabs.currentIndex()
         
         if current_tab_index == 0:
-            is_color_features = (self.category_combo.currentText() == "Color Features")
+            category = self.category_combo.currentText()
+            is_color_features = (category == "Color Features")
+            is_transformer = (category == "Transformer Model")
+            is_community = (category == "Community Model")  # Check for Community Model
+        
+        # Disable feature mode for Color Features, Transformer Models, and Community Models
+        self.feature_mode_combo.setEnabled(not (is_color_features or is_transformer or is_community))
         
-        self.feature_mode_combo.setEnabled(not is_color_features)
+        # If disabled categories are selected, force the combo to "Embed Features"
+        if is_color_features or is_transformer or is_community:
+            self.feature_mode_combo.setCurrentText("Embed Features")
         
         if is_color_features:
             self.feature_mode_combo.setToolTip("Feature Mode is not applicable for Color Features.")
+        elif is_transformer:
+            self.feature_mode_combo.setToolTip("Transformer models always output embedding features.")
+        elif is_community:
+            self.feature_mode_combo.setToolTip("Community models always output embedding features.")
         else:
             self.feature_mode_combo.setToolTip(
                 "Choose 'Predictions' for class probabilities (for uncertainty analysis)\n"
@@ -589,10 +654,13 @@ class ModelSettingsWidget(QGroupBox):
         current_tab_index = self.tabs.currentIndex()
         model_name = ""
         
+        # Standard Model Tab
         if current_tab_index == 0:
             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()
@@ -601,16 +669,34 @@ class ModelSettingsWidget(QGroupBox):
                 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"
+                # Create the display name format to look up in YOLO_MODELS
+                display_name = f"{family_text} ({size_text})"
+                if display_name in self.yolo_models:
+                    model_name = self.yolo_models[display_name]
+                else:
+                    # Fallback to the old method if not found
+                    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()
+                model_display_name = self.community_combo.currentText()
+                # Use the path (value) instead of the name (key) for community models
+                if self.community_configs and model_display_name in self.community_configs:
+                    model_name = self.community_configs[model_display_name]
+                else:
+                    model_name = model_display_name  # Fallback to using the display name
+                    
+            elif category == "Transformer Model":
+                # Get the HuggingFace model ID from the transformer models map
+                selected_display_name = self.transformer_combo.currentText()
+                model_name = self.transformer_models.get(selected_display_name, selected_display_name)
+        
+        # Custom Model Tab
         elif current_tab_index == 1:
             model_name = self.model_path_edit.text()
         
-        feature_mode = self.feature_mode_combo.currentText() if self.feature_mode_combo.isEnabled() else "N/A"
+        feature_mode = self.feature_mode_combo.currentText()
         return model_name, feature_mode
     
 
@@ -748,4 +834,4 @@ class EmbeddingSettingsWidget(QGroupBox):
             if hasattr(self.explorer_window, '_clear_selections'):
                 self.explorer_window._clear_selections()
 
-            self.explorer_window.run_embedding_pipeline()
+            self.explorer_window.run_embedding_pipeline()