coralnet-toolbox 0.0.67__py2.py3-none-any.whl → 0.0.69__py2.py3-none-any.whl

coralnet_toolbox/Explorer/QtExplorer.py

@@ -1,32 +1,33 @@
+import warnings
+
 import os
+
 import numpy as np
 import torch
-import warnings
 
 from ultralytics import YOLO
 
-from coralnet_toolbox.MachineLearning.Community.cfg import get_available_configs
-
 from coralnet_toolbox.Icons import get_icon
 from coralnet_toolbox.utilities import pixmap_to_numpy
 
-from PyQt5.QtGui import QIcon, QPen, QColor, QPainter, QImage, QBrush, QPainterPath, QPolygonF, QMouseEvent
-from PyQt5.QtCore import Qt, QTimer, QSize, QRect, QRectF, QPointF, pyqtSignal, QSignalBlocker, pyqtSlot
-
+from PyQt5.QtGui import QIcon, QPen, QColor, QPainter, QBrush, QPainterPath, QMouseEvent
+from PyQt5.QtCore import Qt, QTimer, QRect, QRectF, QPointF, pyqtSignal, QSignalBlocker, pyqtSlot
 from PyQt5.QtWidgets import (QVBoxLayout, QHBoxLayout, QGraphicsView, QScrollArea,
-                             QGraphicsScene, QPushButton, QComboBox, QLabel, QWidget, QGridLayout,
-                             QMainWindow, QSplitter, QGroupBox, QFormLayout,
-                             QSpinBox, QGraphicsEllipseItem, QGraphicsItem, QSlider,
-                             QListWidget, QDoubleSpinBox, QApplication, QStyle,
-                             QGraphicsRectItem, QRubberBand, QStyleOptionGraphicsItem,
-                             QTabWidget, QLineEdit, QFileDialog)
-
-from .QtAnnotationDataItem import AnnotationDataItem
-from .QtEmbeddingPointItem import EmbeddingPointItem
-from .QtAnnotationImageWidget import AnnotationImageWidget
-from .QtSettingsWidgets import ModelSettingsWidget
-from .QtSettingsWidgets import EmbeddingSettingsWidget
-from .QtSettingsWidgets import AnnotationSettingsWidget
+                             QGraphicsScene, QPushButton, QComboBox, QLabel, QWidget,
+                             QMainWindow, QSplitter, QGroupBox, QSlider, QMessageBox,
+                             QApplication, QGraphicsRectItem, QRubberBand, QMenu,
+                             QWidgetAction, QToolButton, QAction)
+
+from coralnet_toolbox.Explorer.QtFeatureStore import FeatureStore
+from coralnet_toolbox.Explorer.QtDataItem import AnnotationDataItem
+from coralnet_toolbox.Explorer.QtDataItem import EmbeddingPointItem
+from coralnet_toolbox.Explorer.QtDataItem import AnnotationImageWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import ModelSettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import SimilaritySettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import UncertaintySettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import MislabelSettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import EmbeddingSettingsWidget
+from coralnet_toolbox.Explorer.QtSettingsWidgets import AnnotationSettingsWidget
 
 from coralnet_toolbox.QtProgressBar import ProgressBar
 
@@ -34,13 +35,13 @@ try:
     from sklearn.preprocessing import StandardScaler
     from sklearn.decomposition import PCA
     from sklearn.manifold import TSNE
-    from umap import UMAP  
+    from umap import UMAP
 except ImportError:
     print("Warning: sklearn or umap not installed. Some features may be unavailable.")
     StandardScaler = None
     PCA = None
     TSNE = None
-    UMAP = None  
+    UMAP = None
 
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
@@ -50,60 +51,56 @@ warnings.filterwarnings("ignore", category=DeprecationWarning)
 # Constants
 # ----------------------------------------------------------------------------------------------------------------------
 
-
-POINT_SIZE = 15
 POINT_WIDTH = 3
 
-
 # ----------------------------------------------------------------------------------------------------------------------
 # Viewers
 # ----------------------------------------------------------------------------------------------------------------------
 
 
-class EmbeddingViewer(QWidget):  # Change inheritance to QWidget
-    """Custom QGraphicsView for interactive embedding visualization with zooming, panning, and selection."""
-    
-    # Define signal to report selection changes
-    selection_changed = pyqtSignal(list)  # list of all currently selected annotation IDs
-    reset_view_requested = pyqtSignal()  # Signal to reset the view to fit all points
-    
+class EmbeddingViewer(QWidget):
+    """Custom QGraphicsView for interactive embedding visualization with an isolate mode."""
+    selection_changed = pyqtSignal(list)
+    reset_view_requested = pyqtSignal()
+    find_mislabels_requested = pyqtSignal()
+    mislabel_parameters_changed = pyqtSignal(dict) 
+    find_uncertain_requested = pyqtSignal()
+    uncertainty_parameters_changed = pyqtSignal(dict)
+
     def __init__(self, parent=None):
-        # Create the graphics scene first
-        self.graphics_scene = QGraphicsScene()
-        self.graphics_scene.setSceneRect(-5000, -5000, 10000, 10000)
-        
-        # Initialize as a QWidget
+        """Initialize the EmbeddingViewer widget."""
         super(EmbeddingViewer, self).__init__(parent)
         self.explorer_window = parent
-        
-        # Create the actual graphics view
+
+        self.graphics_scene = QGraphicsScene()
+        self.graphics_scene.setSceneRect(-5000, -5000, 10000, 10000)
+
         self.graphics_view = QGraphicsView(self.graphics_scene)
         self.graphics_view.setRenderHint(QPainter.Antialiasing)
         self.graphics_view.setDragMode(QGraphicsView.ScrollHandDrag)
         self.graphics_view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
         self.graphics_view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
         self.graphics_view.setMinimumHeight(200)
-        
-        # Custom rubber_band state variables
+
         self.rubber_band = None
         self.rubber_band_origin = QPointF()
         self.selection_at_press = None
+        self.points_by_id = {}
+        self.previous_selection_ids = set()
+
+        # State for isolate mode
+        self.isolated_mode = False
+        self.isolated_points = set()
         
-        self.points_by_id = {}  # Map annotation ID to embedding point
-        self.previous_selection_ids = set()  # Track previous selection to detect changes
-    
+        self.is_uncertainty_analysis_available = False
+
         self.animation_offset = 0
         self.animation_timer = QTimer()
         self.animation_timer.timeout.connect(self.animate_selection)
         self.animation_timer.setInterval(100)
-        
-        # Connect the scene's selection signal
+
         self.graphics_scene.selectionChanged.connect(self.on_selection_changed)
-        
-        # Setup the UI with header
         self.setup_ui()
-        
-        # Connect mouse events to the graphics view
         self.graphics_view.mousePressEvent = self.mousePressEvent
         self.graphics_view.mouseDoubleClickEvent = self.mouseDoubleClickEvent
         self.graphics_view.mouseReleaseEvent = self.mouseReleaseEvent
@@ -111,36 +108,164 @@ class EmbeddingViewer(QWidget):  # Change inheritance to QWidget
         self.graphics_view.wheelEvent = self.wheelEvent
 
     def setup_ui(self):
-        """Set up the UI with header layout and graphics view."""
+        """Set up the UI with toolbar layout and graphics view."""
         layout = QVBoxLayout(self)
         layout.setContentsMargins(0, 0, 0, 0)
+
+        toolbar_layout = QHBoxLayout()
+
+        # Isolate/Show All buttons
+        self.isolate_button = QPushButton("Isolate Selection")
+        self.isolate_button.setToolTip("Hide all non-selected points")
+        self.isolate_button.clicked.connect(self.isolate_selection)
+        toolbar_layout.addWidget(self.isolate_button)
+
+        self.show_all_button = QPushButton("Show All")
+        self.show_all_button.setToolTip("Show all embedding points")
+        self.show_all_button.clicked.connect(self.show_all_points)
+        toolbar_layout.addWidget(self.show_all_button)
         
-        # Header layout
-        header_layout = QHBoxLayout()
+        toolbar_layout.addWidget(self._create_separator())
+                
+        # Create a QToolButton to have both a primary action and a dropdown menu
+        self.find_mislabels_button = QToolButton()
+        self.find_mislabels_button.setText("Find Potential Mislabels")
+        self.find_mislabels_button.setPopupMode(QToolButton.MenuButtonPopup) # Key change for split-button style
+        self.find_mislabels_button.setToolButtonStyle(Qt.ToolButtonTextOnly)
+        self.find_mislabels_button.setStyleSheet(
+            "QToolButton::menu-indicator {"
+            " subcontrol-position: right center;"
+            " subcontrol-origin: padding;"
+            " left: -4px;"
+            " }"
+        )
+
+        # The primary action (clicking the button) triggers the analysis
+        run_analysis_action = QAction("Find Potential Mislabels", self)
+        run_analysis_action.triggered.connect(self.find_mislabels_requested.emit)
+        self.find_mislabels_button.setDefaultAction(run_analysis_action)
+
+        # The dropdown menu contains the settings
+        mislabel_settings_widget = MislabelSettingsWidget()
+        settings_menu = QMenu(self)
+        widget_action = QWidgetAction(settings_menu)
+        widget_action.setDefaultWidget(mislabel_settings_widget)
+        settings_menu.addAction(widget_action)
+        self.find_mislabels_button.setMenu(settings_menu)
+        
+        # Connect the widget's signal to the viewer's signal
+        mislabel_settings_widget.parameters_changed.connect(self.mislabel_parameters_changed.emit)
+        toolbar_layout.addWidget(self.find_mislabels_button)
+        
+        # Create a QToolButton for uncertainty analysis
+        self.find_uncertain_button = QToolButton()
+        self.find_uncertain_button.setText("Review Uncertain")
+        self.find_uncertain_button.setToolTip(
+            "Find annotations where the model is least confident.\n"
+            "Requires a .pt classification model and 'Predictions' mode."
+        )
+        self.find_uncertain_button.setPopupMode(QToolButton.MenuButtonPopup)
+        self.find_uncertain_button.setToolButtonStyle(Qt.ToolButtonTextOnly)
+        self.find_uncertain_button.setStyleSheet(
+            "QToolButton::menu-indicator { "
+            "subcontrol-position: right center; "
+            "subcontrol-origin: padding; "
+            "left: -4px; }"
+        )
         
-        # Home button
-        self.home_button = QPushButton("Home")
-        self.home_button.setToolTip("Reset view to fit all points")
-        self.home_button.clicked.connect(self.reset_view)
-        header_layout.addWidget(self.home_button)
+        run_uncertainty_action = QAction("Review Uncertain", self)
+        run_uncertainty_action.triggered.connect(self.find_uncertain_requested.emit)
+        self.find_uncertain_button.setDefaultAction(run_uncertainty_action)
+
+        uncertainty_settings_widget = UncertaintySettingsWidget()
+        uncertainty_menu = QMenu(self)
+        uncertainty_widget_action = QWidgetAction(uncertainty_menu)
+        uncertainty_widget_action.setDefaultWidget(uncertainty_settings_widget)
+        uncertainty_menu.addAction(uncertainty_widget_action)
+        self.find_uncertain_button.setMenu(uncertainty_menu)
         
-        # Add stretch to push future controls to the right if needed
-        header_layout.addStretch()
+        uncertainty_settings_widget.parameters_changed.connect(self.uncertainty_parameters_changed.emit)
+        toolbar_layout.addWidget(self.find_uncertain_button)
+
+        toolbar_layout.addStretch()
         
-        layout.addLayout(header_layout)
+        # Home button to reset view
+        self.home_button = QPushButton()
+        self.home_button.setIcon(get_icon("home.png"))
+        self.home_button.setToolTip("Reset view to fit all points")
+        self.home_button.clicked.connect(self.reset_view)
+        toolbar_layout.addWidget(self.home_button)
         
-        # Add the graphics view
+        layout.addLayout(toolbar_layout)
         layout.addWidget(self.graphics_view)
-        # Add a placeholder label when no embedding is available
+
         self.placeholder_label = QLabel(
             "No embedding data available.\nPress 'Apply Embedding' to generate visualization."
         )
         self.placeholder_label.setAlignment(Qt.AlignCenter)
         self.placeholder_label.setStyleSheet("color: gray; font-size: 14px;")
         layout.addWidget(self.placeholder_label)
-        
-        # Initially show placeholder
+
         self.show_placeholder()
+        self._update_toolbar_state()
+        
+    def _create_separator(self):
+        """Creates a vertical separator for the toolbar."""
+        separator = QLabel("|")
+        separator.setStyleSheet("color: gray; margin: 0 5px;")
+        return separator
+        
+    @pyqtSlot()
+    def isolate_selection(self):
+        """Hides all points that are not currently selected."""
+        selected_items = self.graphics_scene.selectedItems()
+        if not selected_items or self.isolated_mode:
+            return
+
+        self.isolated_points = set(selected_items)
+        self.graphics_view.setUpdatesEnabled(False)
+        try:
+            for point in self.points_by_id.values():
+                if point not in self.isolated_points:
+                    point.hide()
+            self.isolated_mode = True
+        finally:
+            self.graphics_view.setUpdatesEnabled(True)
+
+        self._update_toolbar_state()
+
+    @pyqtSlot()
+    def show_all_points(self):
+        """Shows all embedding points, exiting isolated mode."""
+        if not self.isolated_mode:
+            return
+
+        self.isolated_mode = False
+        self.isolated_points.clear()
+        self.graphics_view.setUpdatesEnabled(False)
+        try:
+            for point in self.points_by_id.values():
+                point.show()
+        finally:
+            self.graphics_view.setUpdatesEnabled(True)
+
+        self._update_toolbar_state()
+
+    def _update_toolbar_state(self):
+        """Updates toolbar buttons based on selection and isolation mode."""
+        selection_exists = bool(self.graphics_scene.selectedItems())
+        points_exist = bool(self.points_by_id)
+
+        self.find_mislabels_button.setEnabled(points_exist)
+        self.find_uncertain_button.setEnabled(points_exist and self.is_uncertainty_analysis_available)
+
+        if self.isolated_mode:
+            self.isolate_button.hide()
+            self.show_all_button.show()
+        else:
+            self.isolate_button.show()
+            self.show_all_button.hide()
+            self.isolate_button.setEnabled(selection_exists)
 
     def reset_view(self):
         """Reset the view to fit all embedding points."""
@@ -151,50 +276,87 @@ class EmbeddingViewer(QWidget):  # Change inheritance to QWidget
         self.graphics_view.setVisible(False)
         self.placeholder_label.setVisible(True)
         self.home_button.setEnabled(False)
+        self.find_mislabels_button.setEnabled(False)
+        self.find_uncertain_button.setEnabled(False)
+
+        self.isolate_button.show()
+        self.isolate_button.setEnabled(False)
+        self.show_all_button.hide()
 
     def show_embedding(self):
         """Show the graphics view and hide the placeholder message."""
         self.graphics_view.setVisible(True)
         self.placeholder_label.setVisible(False)
         self.home_button.setEnabled(True)
+        self._update_toolbar_state()
 
     # Delegate graphics view methods
     def setRenderHint(self, hint):
+        """Set render hint for the graphics view."""
         self.graphics_view.setRenderHint(hint)
-    
+
     def setDragMode(self, mode):
+        """Set drag mode for the graphics view."""
         self.graphics_view.setDragMode(mode)
-    
+
     def setTransformationAnchor(self, anchor):
+        """Set transformation anchor for the graphics view."""
         self.graphics_view.setTransformationAnchor(anchor)
-    
+
     def setResizeAnchor(self, anchor):
+        """Set resize anchor for the graphics view."""
         self.graphics_view.setResizeAnchor(anchor)
-    
+
     def mapToScene(self, point):
+        """Map a point to the scene coordinates."""
         return self.graphics_view.mapToScene(point)
-    
+
     def scale(self, sx, sy):
+        """Scale the graphics view."""
         self.graphics_view.scale(sx, sy)
-    
+
     def translate(self, dx, dy):
+        """Translate the graphics view."""
         self.graphics_view.translate(dx, dy)
-    
+
     def fitInView(self, rect, aspect_ratio):
+        """Fit the view to a rectangle with aspect ratio."""
         self.graphics_view.fitInView(rect, aspect_ratio)
 
+    def keyPressEvent(self, event):
+        """Handles key presses for deleting selected points."""
+        if event.key() in (Qt.Key_Delete, Qt.Key_Backspace) and event.modifiers() == Qt.ControlModifier:
+            selected_items = self.graphics_scene.selectedItems()
+            if not selected_items:
+                super().keyPressEvent(event)
+                return
+
+            # Extract the central data items from the selected graphics points
+            data_items_to_delete = [
+                item.data_item for item in selected_items if isinstance(item, EmbeddingPointItem)
+            ]
+
+            # Delegate the actual deletion to the main ExplorerWindow
+            if data_items_to_delete:
+                self.explorer_window.delete_data_items(data_items_to_delete)
+
+            event.accept()
+        else:
+            super().keyPressEvent(event)
+
     def mousePressEvent(self, event):
         """Handle mouse press for selection (point or rubber band) and panning."""
         if event.button() == Qt.LeftButton and event.modifiers() == Qt.ControlModifier:
-            # Check if the click is on an existing point
             item_at_pos = self.graphics_view.itemAt(event.pos())
             if isinstance(item_at_pos, EmbeddingPointItem):
-                # If so, toggle its selection state and do nothing else
                 self.graphics_view.setDragMode(QGraphicsView.NoDrag)
-                item_at_pos.setSelected(not item_at_pos.isSelected())
-                return  # Event handled
+                # The viewer (controller) directly changes the state on the data item.
+                is_currently_selected = item_at_pos.data_item.is_selected
+                item_at_pos.data_item.set_selected(not is_currently_selected)
+                item_at_pos.setSelected(not is_currently_selected)  # Keep scene selection in sync
+                self.on_selection_changed()  # Manually trigger update
+                return
 
-            # If the click was on the background, proceed with rubber band selection
             self.selection_at_press = set(self.graphics_scene.selectedItems())
             self.graphics_view.setDragMode(QGraphicsView.NoDrag)
             self.rubber_band_origin = self.graphics_view.mapToScene(event.pos())
@@ -204,210 +366,186 @@ class EmbeddingViewer(QWidget):  # Change inheritance to QWidget
             self.graphics_scene.addItem(self.rubber_band)
 
         elif event.button() == Qt.RightButton:
-            # Handle panning
             self.graphics_view.setDragMode(QGraphicsView.ScrollHandDrag)
             left_event = QMouseEvent(event.type(), event.localPos(), Qt.LeftButton, Qt.LeftButton, event.modifiers())
             QGraphicsView.mousePressEvent(self.graphics_view, left_event)
         else:
-            # Handle standard single-item selection
             self.graphics_view.setDragMode(QGraphicsView.NoDrag)
             QGraphicsView.mousePressEvent(self.graphics_view, event)
-            
+
     def mouseDoubleClickEvent(self, event):
         """Handle double-click to clear selection and reset the main view."""
         if event.button() == Qt.LeftButton:
-            # Clear selection if any items are selected
             if self.graphics_scene.selectedItems():
-                self.graphics_scene.clearSelection()  # This triggers on_selection_changed
-            
-            # Signal the main window to revert from isolation mode
+                self.graphics_scene.clearSelection()
             self.reset_view_requested.emit()
             event.accept()
         else:
-            # Pass other double-clicks to the base class
             super().mouseDoubleClickEvent(event)
 
     def mouseMoveEvent(self, event):
         """Handle mouse move for dynamic selection and panning."""
         if self.rubber_band:
-            # Update the rubber band geometry
+            # Update the rubber band rectangle as the mouse moves
             current_pos = self.graphics_view.mapToScene(event.pos())
             self.rubber_band.setRect(QRectF(self.rubber_band_origin, current_pos).normalized())
-            
+            # Create a selection path from the rubber band rectangle
             path = QPainterPath()
             path.addRect(self.rubber_band.rect())
-
-            # Block signals to perform a compound selection operation
+            # Block signals to avoid recursive selectionChanged events
             self.graphics_scene.blockSignals(True)
-
-            # 1. Perform the "fancy" dynamic selection, which replaces the current selection
-            #    with only the items inside the rubber band.
             self.graphics_scene.setSelectionArea(path)
-            
-            # 2. Add back the items that were selected at the start of the drag.
+            # Restore selection for items that were already selected at press
             if self.selection_at_press:
                 for item in self.selection_at_press:
                     item.setSelected(True)
-            
-            # Unblock signals and manually trigger our handler to process the final result.
             self.graphics_scene.blockSignals(False)
+            # Manually trigger selection changed logic
             self.on_selection_changed()
-
         elif event.buttons() == Qt.RightButton:
-            # Handle right-click panning
-            left_event = QMouseEvent(event.type(), 
-                                     event.localPos(), 
-                                     Qt.LeftButton, 
-                                     Qt.LeftButton, 
-                                     event.modifiers())
+            # Forward right-drag as left-drag for panning
+            left_event = QMouseEvent(event.type(), event.localPos(), Qt.LeftButton, Qt.LeftButton, event.modifiers())
             QGraphicsView.mouseMoveEvent(self.graphics_view, left_event)
         else:
+            # Default mouse move handling
             QGraphicsView.mouseMoveEvent(self.graphics_view, event)
 
     def mouseReleaseEvent(self, event):
         """Handle mouse release to finalize the action and clean up."""
         if self.rubber_band:
-            # Clean up the visual rectangle
             self.graphics_scene.removeItem(self.rubber_band)
             self.rubber_band = None
-
-            # Clean up the stored selection state.
             self.selection_at_press = None
-            
         elif event.button() == Qt.RightButton:
-            # Finalize the pan
-            left_event = QMouseEvent(event.type(), 
-                                     event.localPos(), 
-                                     Qt.LeftButton, 
-                                     Qt.LeftButton, 
-                                     event.modifiers())
+            left_event = QMouseEvent(event.type(), event.localPos(), Qt.LeftButton, Qt.LeftButton, event.modifiers())
             QGraphicsView.mouseReleaseEvent(self.graphics_view, left_event)
             self.graphics_view.setDragMode(QGraphicsView.NoDrag)
         else:
-            # Finalize a single click
             QGraphicsView.mouseReleaseEvent(self.graphics_view, event)
             self.graphics_view.setDragMode(QGraphicsView.NoDrag)
-    
+
     def wheelEvent(self, event):
         """Handle mouse wheel for zooming."""
         zoom_in_factor = 1.25
         zoom_out_factor = 1 / zoom_in_factor
 
+        # Set anchor points so zoom occurs at mouse position
         self.graphics_view.setTransformationAnchor(QGraphicsView.NoAnchor)
         self.graphics_view.setResizeAnchor(QGraphicsView.NoAnchor)
 
+        # Get the scene position before zooming
         old_pos = self.graphics_view.mapToScene(event.pos())
+
+        # Determine zoom direction
         zoom_factor = zoom_in_factor if event.angleDelta().y() > 0 else zoom_out_factor
+
+        # Apply zoom
         self.graphics_view.scale(zoom_factor, zoom_factor)
+
+        # Get the scene position after zooming
         new_pos = self.graphics_view.mapToScene(event.pos())
-        
+
+        # Translate view to keep mouse position stable
         delta = new_pos - old_pos
         self.graphics_view.translate(delta.x(), delta.y())
 
     def update_embeddings(self, data_items):
-        """Update the embedding visualization with new data.
+        """Update the embedding visualization. Creates an EmbeddingPointItem for
+        each AnnotationDataItem and links them."""
+        # Reset isolation state when loading new points
+        if self.isolated_mode:
+            self.show_all_points()
 
-        Args:
-            data_items: List of AnnotationDataItem objects.
-        """
         self.clear_points()
-        
         for item in data_items:
-            point = EmbeddingPointItem(0, 0, POINT_SIZE, POINT_SIZE)
-            point.setPos(item.embedding_x, item.embedding_y)
-            
-            # No need to set initial brush - paint() will handle it
-            point.setPen(QPen(QColor("black"), POINT_WIDTH))
-            
-            point.setFlag(QGraphicsItem.ItemIgnoresTransformations)
-            point.setFlag(QGraphicsItem.ItemIsSelectable)
-            
-            # This is the crucial link: store the shared AnnotationDataItem
-            point.setData(0, item)
-            
+            point = EmbeddingPointItem(item)
             self.graphics_scene.addItem(point)
             self.points_by_id[item.annotation.id] = point
-            
+        
+        # Ensure buttons are in the correct initial state
+        self._update_toolbar_state()
+
     def clear_points(self):
         """Clear all embedding points from the scene."""
+        if self.isolated_mode:
+            self.show_all_points()
+
         for point in self.points_by_id.values():
             self.graphics_scene.removeItem(point)
         self.points_by_id.clear()
+        self._update_toolbar_state()
 
     def on_selection_changed(self):
-        """Handle point selection changes and emit a signal to the controller."""
-        # Check if graphics_scene still exists and is valid
-        if not self.graphics_scene or not hasattr(self.graphics_scene, 'selectedItems'):
+        """
+        Handles selection changes in the scene. Updates the central data model
+        and emits a signal to notify other parts of the application.
+        """
+        if not self.graphics_scene:
             return
-            
         try:
             selected_items = self.graphics_scene.selectedItems()
         except RuntimeError:
-            # Scene has been deleted
             return
-            
-        current_selection_ids = {item.data(0).annotation.id for item in selected_items}
 
-        # If the selection has actually changed, update the model and emit
+        current_selection_ids = {item.data_item.annotation.id for item in selected_items}
+
         if current_selection_ids != self.previous_selection_ids:
-            # Update the central model (the AnnotationDataItem) for all points
             for point_id, point in self.points_by_id.items():
                 is_selected = point_id in current_selection_ids
-                point.data(0).set_selected(is_selected)
+                point.data_item.set_selected(is_selected)
 
-            # Emit the complete list of currently selected IDs
             self.selection_changed.emit(list(current_selection_ids))
             self.previous_selection_ids = current_selection_ids
 
-        # Handle local animation - check if animation_timer still exists
         if hasattr(self, 'animation_timer') and self.animation_timer:
             self.animation_timer.stop()
-            
+
         for point in self.points_by_id.values():
             if not point.isSelected():
                 point.setPen(QPen(QColor("black"), POINT_WIDTH))
-        
         if selected_items and hasattr(self, 'animation_timer') and self.animation_timer:
             self.animation_timer.start()
 
+        # Update button states based on new selection
+        self._update_toolbar_state()
+
     def animate_selection(self):
-        """Animate selected points with marching ants effect using darker versions of point colors."""
-        # Check if graphics_scene still exists and is valid
-        if not self.graphics_scene or not hasattr(self.graphics_scene, 'selectedItems'):
+        """Animate selected points with a marching ants effect."""
+        if not self.graphics_scene:
             return
-            
         try:
             selected_items = self.graphics_scene.selectedItems()
         except RuntimeError:
-            # Scene has been deleted
             return
-            
+
         self.animation_offset = (self.animation_offset + 1) % 20
-        
-        # This logic remains the same. It applies the custom pen to the selected items.
-        # Because the items are EmbeddingPointItem, the default selection box won't be drawn.
         for item in selected_items:
-            original_color = item.brush().color()
+            # Get the color directly from the source of truth
+            original_color = item.data_item.effective_color
             darker_color = original_color.darker(150)
-
             animated_pen = QPen(darker_color, POINT_WIDTH)
             animated_pen.setStyle(Qt.CustomDashLine)
             animated_pen.setDashPattern([1, 2])
             animated_pen.setDashOffset(self.animation_offset)
-            
             item.setPen(animated_pen)
-            
+
     def render_selection_from_ids(self, selected_ids):
-        """Update the visual selection of points based on a set of IDs from the controller."""
-        # Block this scene's own selectionChanged signal to prevent an infinite loop
+        """
+        Updates the visual selection of points based on a set of annotation IDs
+        provided by an external controller.
+        """
         blocker = QSignalBlocker(self.graphics_scene)
-        
+
         for ann_id, point in self.points_by_id.items():
-            point.setSelected(ann_id in selected_ids)
-            
-        self.previous_selection_ids = selected_ids
-        
-        # Trigger animation update
+            is_selected = ann_id in selected_ids
+            # 1. Update the state on the central data item
+            point.data_item.set_selected(is_selected)
+            # 2. Update the selection state of the graphics item itself
+            point.setSelected(is_selected)
+
+        blocker.unblock()
+
+        # Manually trigger on_selection_changed to update animation and emit signals
         self.on_selection_changed()
 
     def fit_view_to_points(self):
@@ -415,40 +553,40 @@ class EmbeddingViewer(QWidget):  # Change inheritance to QWidget
         if self.points_by_id:
             self.graphics_view.fitInView(self.graphics_scene.itemsBoundingRect(), Qt.KeepAspectRatio)
         else:
-            # If no points, reset to default view
             self.graphics_view.fitInView(-2500, -2500, 5000, 5000, Qt.KeepAspectRatio)
-            
+
 
 class AnnotationViewer(QScrollArea):
     """Scrollable grid widget for displaying annotation image crops with selection,
-    filtering, and isolation support.
-    """
-    
-    # Define signals to report changes to the ExplorerWindow
-    selection_changed = pyqtSignal(list)  # list of changed annotation IDs
-    preview_changed = pyqtSignal(list)   # list of annotation IDs with new previews
-    reset_view_requested = pyqtSignal()  # Signal to reset the view to fit all points
+    filtering, and isolation support. Acts as a controller for the widgets."""
+    selection_changed = pyqtSignal(list)
+    preview_changed = pyqtSignal(list)
+    reset_view_requested = pyqtSignal()
+    find_similar_requested = pyqtSignal()
 
     def __init__(self, parent=None):
+        """Initialize the AnnotationViewer widget."""
         super(AnnotationViewer, self).__init__(parent)
+        self.explorer_window = parent
+
         self.annotation_widgets_by_id = {}
         self.selected_widgets = []
         self.last_selected_index = -1
         self.current_widget_size = 96
-        
         self.selection_at_press = set()
         self.rubber_band = None
         self.rubber_band_origin = None
         self.drag_threshold = 5
         self.mouse_pressed_on_widget = False
-        
         self.preview_label_assignments = {}
         self.original_label_assignments = {}
-        
-        # New state variables for Isolate/Focus mode
         self.isolated_mode = False
         self.isolated_widgets = set()
 
+        # State for new sorting options
+        self.active_ordered_ids = []
+        self.is_confidence_sort_available = False
+
         self.setup_ui()
 
     def setup_ui(self):
@@ -457,49 +595,61 @@ class AnnotationViewer(QScrollArea):
         self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
         self.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
 
-        # Main container and layout
         main_container = QWidget()
         main_layout = QVBoxLayout(main_container)
         main_layout.setContentsMargins(0, 0, 0, 0)
-        main_layout.setSpacing(4)  # Add a little space between toolbar and content
+        main_layout.setSpacing(4)
 
-        # --- New Toolbar ---
         toolbar_widget = QWidget()
         toolbar_layout = QHBoxLayout(toolbar_widget)
         toolbar_layout.setContentsMargins(4, 2, 4, 2)
 
-        # Isolate/Focus controls
         self.isolate_button = QPushButton("Isolate Selection")
-        isolate_icon = get_icon("focus.png")
-        if not isolate_icon.isNull():
-            self.isolate_button.setIcon(isolate_icon)
         self.isolate_button.setToolTip("Hide all non-selected annotations")
         self.isolate_button.clicked.connect(self.isolate_selection)
         toolbar_layout.addWidget(self.isolate_button)
 
         self.show_all_button = QPushButton("Show All")
-        show_all_icon = get_icon("show_all.png")
-        if not show_all_icon.isNull():
-            self.show_all_button.setIcon(show_all_icon)
         self.show_all_button.setToolTip("Show all filtered annotations")
         self.show_all_button.clicked.connect(self.show_all_annotations)
         toolbar_layout.addWidget(self.show_all_button)
 
-        # Add a separator
         toolbar_layout.addWidget(self._create_separator())
 
-        # Sort controls
         sort_label = QLabel("Sort By:")
         toolbar_layout.addWidget(sort_label)
         self.sort_combo = QComboBox()
-        self.sort_combo.addItems(["None", "Label", "Image"])
+        # Remove "Similarity" as it's now an implicit action
+        self.sort_combo.addItems(["None", "Label", "Image", "Confidence"])
+        self.sort_combo.insertSeparator(3)  # Add separator before "Confidence"
         self.sort_combo.currentTextChanged.connect(self.on_sort_changed)
         toolbar_layout.addWidget(self.sort_combo)
+        
+        toolbar_layout.addWidget(self._create_separator())
+        
+        self.find_similar_button = QToolButton()
+        self.find_similar_button.setText("Find Similar")
+        self.find_similar_button.setToolTip("Find annotations visually similar to the selection.")
+        self.find_similar_button.setPopupMode(QToolButton.MenuButtonPopup)
+        self.find_similar_button.setToolButtonStyle(Qt.ToolButtonTextOnly)
+        self.find_similar_button.setStyleSheet(
+            "QToolButton::menu-indicator { subcontrol-position: right center; subcontrol-origin: padding; left: -4px; }"
+        )
 
-        # Add a spacer to push the size controls to the right
+        run_similar_action = QAction("Find Similar", self)
+        run_similar_action.triggered.connect(self.find_similar_requested.emit)
+        self.find_similar_button.setDefaultAction(run_similar_action)
+        
+        self.similarity_settings_widget = SimilaritySettingsWidget()
+        settings_menu = QMenu(self)
+        widget_action = QWidgetAction(settings_menu)
+        widget_action.setDefaultWidget(self.similarity_settings_widget)
+        settings_menu.addAction(widget_action)
+        self.find_similar_button.setMenu(settings_menu)
+        toolbar_layout.addWidget(self.find_similar_button)
+        
         toolbar_layout.addStretch()
 
-        # Size controls
         size_label = QLabel("Size:")
         toolbar_layout.addWidget(size_label)
         self.size_slider = QSlider(Qt.Horizontal)
@@ -514,23 +664,63 @@ class AnnotationViewer(QScrollArea):
         self.size_value_label = QLabel("96")
         self.size_value_label.setMinimumWidth(30)
         toolbar_layout.addWidget(self.size_value_label)
-
         main_layout.addWidget(toolbar_widget)
-        
-        # --- Content Area ---
+
         self.content_widget = QWidget()
         content_scroll = QScrollArea()
         content_scroll.setWidgetResizable(True)
         content_scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
         content_scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
         content_scroll.setWidget(self.content_widget)
-        
+
         main_layout.addWidget(content_scroll)
         self.setWidget(main_container)
 
-        # Set the initial state of the toolbar buttons
+        # Set the initial state of the sort options
+        self._update_sort_options_state()
         self._update_toolbar_state()
-            
+        
+    def _create_separator(self):
+        """Creates a vertical separator for the toolbar."""
+        separator = QLabel("|")
+        separator.setStyleSheet("color: gray; margin: 0 5px;")
+        return separator
+    
+    def _update_sort_options_state(self):
+        """Enable/disable sort options based on available data."""
+        model = self.sort_combo.model()
+
+        # Enable/disable "Confidence" option
+        confidence_item_index = self.sort_combo.findText("Confidence")
+        if confidence_item_index != -1:
+            model.item(confidence_item_index).setEnabled(self.is_confidence_sort_available)
+
+    def handle_annotation_context_menu(self, widget, event):
+        """Handle context menu requests (e.g., right-click) on an annotation widget."""
+        if event.modifiers() == Qt.ControlModifier:
+            explorer = self.explorer_window
+            image_path = widget.annotation.image_path
+            annotation_to_select = widget.annotation
+
+            if hasattr(explorer, 'annotation_window'):
+                # Check if the image needs to be changed
+                if explorer.annotation_window.current_image_path != image_path:
+                    if hasattr(explorer.annotation_window, 'set_image'):
+                        explorer.annotation_window.set_image(image_path)
+
+                # Now, select the annotation in the annotation_window
+                if hasattr(explorer.annotation_window, 'select_annotation'):
+                    # This method by default unselects other annotations
+                    explorer.annotation_window.select_annotation(annotation_to_select)
+
+                # Also clear any existing selection in the explorer window itself
+                explorer.annotation_viewer.clear_selection()
+                explorer.embedding_viewer.render_selection_from_ids(set())
+                explorer.update_label_window_selection()
+                explorer.update_button_states()
+
+            event.accept()
+
     @pyqtSlot()
     def isolate_selection(self):
         """Hides all annotation widgets that are not currently selected."""
@@ -549,32 +739,64 @@ class AnnotationViewer(QScrollArea):
             self.content_widget.setUpdatesEnabled(True)
 
         self._update_toolbar_state()
+        self.explorer_window.main_window.label_window.update_annotation_count()
+
+    def display_and_isolate_ordered_results(self, ordered_ids):
+        """
+        Isolates the view to a specific set of ordered widgets, ensuring the 
+        grid is always updated. This is the new primary method for showing 
+        similarity results.
+        """
+        self.active_ordered_ids = ordered_ids
+        
+        # Render the selection based on the new order
+        self.render_selection_from_ids(set(ordered_ids)) 
+
+        # Now, perform the isolation logic directly to bypass the guard clause
+        self.isolated_widgets = set(self.selected_widgets)
+        self.content_widget.setUpdatesEnabled(False)
+        try:
+            for widget in self.annotation_widgets_by_id.values():
+                # Show widget if it's in our target set, hide otherwise
+                if widget in self.isolated_widgets:
+                    widget.show()
+                else:
+                    widget.hide()
+                    
+            self.isolated_mode = True
+            self.recalculate_widget_positions()  # Crucial grid update
+        finally:
+            self.content_widget.setUpdatesEnabled(True)
+
+        self._update_toolbar_state()
+        self.explorer_window.main_window.label_window.update_annotation_count()
 
     @pyqtSlot()
     def show_all_annotations(self):
         """Shows all annotation widgets, exiting the isolated mode."""
         if not self.isolated_mode:
             return
-            
+
         self.isolated_mode = False
         self.isolated_widgets.clear()
-        
+        self.active_ordered_ids = []  # Clear similarity sort context
+
         self.content_widget.setUpdatesEnabled(False)
         try:
+            # Show all widgets that are managed by the viewer
             for widget in self.annotation_widgets_by_id.values():
                 widget.show()
+
             self.recalculate_widget_positions()
         finally:
             self.content_widget.setUpdatesEnabled(True)
-            
+
         self._update_toolbar_state()
-        
+        self.explorer_window.main_window.label_window.update_annotation_count()
+
     def _update_toolbar_state(self):
-        """Updates the visibility and enabled state of the toolbar buttons
-        based on the current selection and isolation mode.
-        """
+        """Updates the toolbar buttons based on selection and isolation mode."""
         selection_exists = bool(self.selected_widgets)
-        
         if self.isolated_mode:
             self.isolate_button.hide()
             self.show_all_button.show()
@@ -583,65 +805,70 @@ class AnnotationViewer(QScrollArea):
             self.isolate_button.show()
             self.show_all_button.hide()
             self.isolate_button.setEnabled(selection_exists)
-            
-    def _create_separator(self):
-        """Create a vertical separator line for the toolbar."""
-        separator = QLabel("|")
-        separator.setStyleSheet("color: gray; margin: 0 5px;")
-        return separator
 
     def on_sort_changed(self, sort_type):
         """Handle sort type change."""
+        self.active_ordered_ids = []  # Clear any special ordering
         self.recalculate_widget_positions()
 
+    def set_confidence_sort_availability(self, is_available):
+        """Sets the availability of the confidence sort option."""
+        self.is_confidence_sort_available = is_available
+        self._update_sort_options_state()
+
     def _get_sorted_widgets(self):
         """Get widgets sorted according to the current sort setting."""
+        # If a specific order is active (e.g., from similarity search), use it.
+        if self.active_ordered_ids:
+            widget_map = {w.data_item.annotation.id: w for w in self.annotation_widgets_by_id.values()}
+            ordered_widgets = [widget_map[ann_id] for ann_id in self.active_ordered_ids if ann_id in widget_map]
+            return ordered_widgets
+
+        # Otherwise, use the dropdown sort logic
         sort_type = self.sort_combo.currentText()
-        
-        if sort_type == "None":
-            return list(self.annotation_widgets_by_id.values())
-        
         widgets = list(self.annotation_widgets_by_id.values())
-        
+
         if sort_type == "Label":
             widgets.sort(key=lambda w: w.data_item.effective_label.short_label_code)
         elif sort_type == "Image":
             widgets.sort(key=lambda w: os.path.basename(w.data_item.annotation.image_path))
+        elif sort_type == "Confidence":
+            # Sort by confidence, descending. Handles cases with no confidence gracefully.
+            widgets.sort(key=lambda w: w.data_item.get_effective_confidence(), reverse=True)
         
         return widgets
 
     def _group_widgets_by_sort_key(self, widgets):
-        """Group widgets by the current sort key and return groups with headers."""
+        """Group widgets by the current sort key."""
         sort_type = self.sort_combo.currentText()
-        
-        if sort_type == "None":
+        if not self.active_ordered_ids and sort_type == "None":
             return [("", widgets)]
         
+        if self.active_ordered_ids: # Don't show group headers for similarity results
+            return [("", widgets)]
+
         groups = []
         current_group = []
         current_key = None
-        
         for widget in widgets:
             if sort_type == "Label":
                 key = widget.data_item.effective_label.short_label_code
             elif sort_type == "Image":
                 key = os.path.basename(widget.data_item.annotation.image_path)
             else:
-                key = ""
+                key = "" # No headers for Confidence or None
             
-            if current_key != key:
+            if key and current_key != key:
                 if current_group:
                     groups.append((current_key, current_group))
                 current_group = [widget]
                 current_key = key
             else:
                 current_group.append(widget)
-        
         if current_group:
             groups.append((current_key, current_group))
-        
         return groups
-    
+
     def _clear_separator_labels(self):
         """Remove any existing group header labels."""
         if hasattr(self, '_group_headers'):
@@ -654,25 +881,22 @@ class AnnotationViewer(QScrollArea):
         """Create a group header label."""
         if not hasattr(self, '_group_headers'):
             self._group_headers = []
-        
-        header = QLabel(text)
-        header.setParent(self.content_widget)
-        header.setStyleSheet("""
-            QLabel {
-                font-weight: bold;
-                font-size: 12px;
-                color: #555;
-                background-color: #f0f0f0;
-                border: 1px solid #ccc;
-                border-radius: 3px;
-                padding: 5px 8px;
-                margin: 2px 0px;
-            }
-        """)
-        header.setFixedHeight(30)  # Increased from 25 to 30
+        header = QLabel(text, self.content_widget)
+        header.setStyleSheet(
+            "QLabel {"
+            " font-weight: bold;"
+            " font-size: 12px;"
+            " color: #555;"
+            " background-color: #f0f0f0;"
+            " border: 1px solid #ccc;"
+            " border-radius: 3px;"
+            " padding: 5px 8px;"
+            " margin: 2px 0px;"
+            " }"
+        )
+        header.setFixedHeight(30)
         header.setMinimumWidth(self.viewport().width() - 20)
         header.show()
-        
         self._group_headers.append(header)
         return header
 
@@ -680,298 +904,251 @@ class AnnotationViewer(QScrollArea):
         """Handle slider value change to resize annotation widgets."""
         if value % 2 != 0:
             value -= 1
+
         self.current_widget_size = value
         self.size_value_label.setText(str(value))
-        
-        # Disable updates for performance while resizing many items
         self.content_widget.setUpdatesEnabled(False)
+
         for widget in self.annotation_widgets_by_id.values():
-            widget.update_height(value)  # Call the new, more descriptive method
-        self.content_widget.setUpdatesEnabled(True)
+            widget.update_height(value)
 
-        # After resizing, reflow the layout
+        self.content_widget.setUpdatesEnabled(True)
         self.recalculate_widget_positions()
 
-    def recalculate_grid_layout(self):
-        """Recalculate the grid layout based on current widget width."""
-        if not self.annotation_widgets_by_id:
-            return
-            
-        available_width = self.viewport().width() - 20
-        widget_width = self.current_widget_size + self.grid_layout.spacing()
-        cols = max(1, available_width // widget_width)
-        
-        for i, widget in enumerate(self.annotation_widgets_by_id.values()):
-            self.grid_layout.addWidget(widget, i // cols, i % cols)
-            
     def recalculate_widget_positions(self):
         """Manually positions widgets in a flow layout with sorting and group headers."""
         if not self.annotation_widgets_by_id:
             self.content_widget.setMinimumSize(1, 1)
             return
 
-        # Clear any existing separator labels
         self._clear_separator_labels()
-
-        # Get sorted widgets
-        all_widgets = self._get_sorted_widgets()
-        
-        # Filter to only visible widgets
-        visible_widgets = [w for w in all_widgets if not w.isHidden()]
-        
+        visible_widgets = [w for w in self._get_sorted_widgets() if not w.isHidden()]
         if not visible_widgets:
             self.content_widget.setMinimumSize(1, 1)
             return
 
-        # Group widgets by sort key
+        # Create groups based on the current sort key
         groups = self._group_widgets_by_sort_key(visible_widgets)
-
-        # Calculate spacing
         spacing = max(5, int(self.current_widget_size * 0.08))
         available_width = self.viewport().width()
-        
         x, y = spacing, spacing
         max_height_in_row = 0
 
+        # Calculate the maximum height of the widgets in each row
         for group_name, group_widgets in groups:
-            # Add group header if sorting is enabled and group has a name
             if group_name and self.sort_combo.currentText() != "None":
-                # Ensure we're at the start of a new line for headers
                 if x > spacing:
                     x = spacing
                     y += max_height_in_row + spacing
                     max_height_in_row = 0
-                
-                # Create and position header label
                 header_label = self._create_group_header(group_name)
                 header_label.move(x, y)
-                
-                # Move to next line after header
                 y += header_label.height() + spacing
                 x = spacing
                 max_height_in_row = 0
 
-            # Position widgets in this group
             for widget in group_widgets:
                 widget_size = widget.size()
-                
-                # Check if widget fits on current line
                 if x > spacing and x + widget_size.width() > available_width:
                     x = spacing
                     y += max_height_in_row + spacing
                     max_height_in_row = 0
-
                 widget.move(x, y)
                 x += widget_size.width() + spacing
                 max_height_in_row = max(max_height_in_row, widget_size.height())
 
-        # Update content widget size
         total_height = y + max_height_in_row + spacing
         self.content_widget.setMinimumSize(available_width, total_height)
-            
+
     def update_annotations(self, data_items):
-        """Update displayed annotations, creating new widgets. This will also
-        reset any active isolation view.
-        """
-        # Reset isolation state before updating to avoid confusion
+        """Update displayed annotations, creating new widgets for them."""
         if self.isolated_mode:
             self.show_all_annotations()
-            
-        # Clear any existing widgets and ensure they are deleted
+
         for widget in self.annotation_widgets_by_id.values():
             widget.setParent(None)
             widget.deleteLater()
-            
+
         self.annotation_widgets_by_id.clear()
         self.selected_widgets.clear()
         self.last_selected_index = -1
 
-        # Create new widgets, parenting them to the content_widget
         for data_item in data_items:
             annotation_widget = AnnotationImageWidget(
-                data_item, 
-                self.current_widget_size,
-                annotation_viewer=self,
-                parent=self.content_widget
-            )
-            annotation_widget.show() 
+                data_item, self.current_widget_size, self, self.content_widget)
+
+            annotation_widget.show()
             self.annotation_widgets_by_id[data_item.annotation.id] = annotation_widget
-        
+
         self.recalculate_widget_positions()
-        # Ensure toolbar is in the correct state after a refresh
         self._update_toolbar_state()
 
     def resizeEvent(self, event):
         """On window resize, reflow the annotation widgets."""
         super().resizeEvent(event)
-        # Use a QTimer to avoid rapid, expensive reflows while dragging the resize handle
         if not hasattr(self, '_resize_timer'):
             self._resize_timer = QTimer(self)
             self._resize_timer.setSingleShot(True)
             self._resize_timer.timeout.connect(self.recalculate_widget_positions)
-        # Restart the timer on each resize event
-        self._resize_timer.start(100)  # 100ms delay
+        self._resize_timer.start(100)
+
+    def keyPressEvent(self, event):
+        """Handles key presses for deleting selected annotations."""
+        if event.key() in (Qt.Key_Delete, Qt.Key_Backspace) and event.modifiers() == Qt.ControlModifier:
+            if not self.selected_widgets:
+                super().keyPressEvent(event)
+                return
+
+            # Extract the central data items from the selected widgets
+            data_items_to_delete = [widget.data_item for widget in self.selected_widgets]
+
+            # Delegate the actual deletion to the main ExplorerWindow
+            if data_items_to_delete:
+                self.explorer_window.delete_data_items(data_items_to_delete)
+
+            event.accept()
+        else:
+            super().keyPressEvent(event)
 
     def mousePressEvent(self, event):
         """Handle mouse press for starting rubber band selection OR clearing selection."""
-        
-        # Handle plain left-clicks
         if event.button() == Qt.LeftButton:
-            
-            # This is the new logic for clearing selection on a background click.
-            if not event.modifiers():  # Check for NO modifiers (e.g., Ctrl, Shift)
-                
+            if not event.modifiers():
+                # If left click with no modifiers, check if click is outside widgets
                 is_on_widget = False
                 child_at_pos = self.childAt(event.pos())
 
-                # Determine if the click was on an actual annotation widget or empty space
                 if child_at_pos:
                     widget = child_at_pos
+                    # Traverse up the parent chain to see if click is on an annotation widget
                     while widget and widget != self:
                         if hasattr(widget, 'annotation_viewer') and widget.annotation_viewer == self:
                             is_on_widget = True
                             break
                         widget = widget.parent()
-                
-                # If click was on empty space AND something is currently selected...
+
+                # If click is outside widgets and there is a selection, clear it
                 if not is_on_widget and self.selected_widgets:
-                    # Get IDs of widgets that are about to be deselected to emit a signal
                     changed_ids = [w.data_item.annotation.id for w in self.selected_widgets]
                     self.clear_selection()
                     self.selection_changed.emit(changed_ids)
-                    # The event is handled, but we don't call super() to prevent
-                    # the scroll area from doing anything else, like starting a drag.
                     return
 
-            # Handle Ctrl+Click for rubber band
             elif event.modifiers() == Qt.ControlModifier:
-                # Store the set of currently selected items.
+                # Start rubber band selection with Ctrl+Left click
                 self.selection_at_press = set(self.selected_widgets)
                 self.rubber_band_origin = event.pos()
-                # We determine mouse_pressed_on_widget here but use it in mouseMove
                 self.mouse_pressed_on_widget = False
                 child_widget = self.childAt(event.pos())
                 if child_widget:
                     widget = child_widget
+                    # Check if click is on a widget to avoid starting rubber band
                     while widget and widget != self:
                         if hasattr(widget, 'annotation_viewer') and widget.annotation_viewer == self:
                             self.mouse_pressed_on_widget = True
                             break
                         widget = widget.parent()
                 return
-                
-        # Handle right-clicks
+
         elif event.button() == Qt.RightButton:
+            # Ignore right clicks
             event.ignore()
             return
-            
-        # For all other cases (e.g., a click on a widget that should be handled
-        # by the widget itself), pass the event to the default handler.
+
+        # Default handler for other cases
         super().mousePressEvent(event)
-        
+
     def mouseDoubleClickEvent(self, event):
         """Handle double-click to clear selection and exit isolation mode."""
         if event.button() == Qt.LeftButton:
             changed_ids = []
-            
-            # If items are selected, clear the selection and record their IDs
             if self.selected_widgets:
                 changed_ids = [w.data_item.annotation.id for w in self.selected_widgets]
                 self.clear_selection()
                 self.selection_changed.emit(changed_ids)
-
-            # If in isolation mode, revert to showing all annotations
             if self.isolated_mode:
                 self.show_all_annotations()
-            
-            # Signal the main window to reset its view (e.g., switch tabs)
             self.reset_view_requested.emit()
             event.accept()
         else:
             super().mouseDoubleClickEvent(event)
-    
+
     def mouseMoveEvent(self, event):
         """Handle mouse move for DYNAMIC rubber band selection."""
-        if self.rubber_band_origin is None or \
-           event.buttons() != Qt.LeftButton or \
-           event.modifiers() != Qt.ControlModifier:
+        # Only proceed if Ctrl+Left mouse drag is active and not on a widget
+        if (
+            self.rubber_band_origin is None or
+            event.buttons() != Qt.LeftButton or
+            event.modifiers() != Qt.ControlModifier
+        ):
             super().mouseMoveEvent(event)
             return
 
-        # If the mouse was pressed on a widget, let that widget handle the event.
         if self.mouse_pressed_on_widget:
+            # If drag started on a widget, do not start rubber band
             super().mouseMoveEvent(event)
             return
 
-        # Only start the rubber band after dragging a minimum distance
+        # Only start selection if drag distance exceeds threshold
         distance = (event.pos() - self.rubber_band_origin).manhattanLength()
         if distance < self.drag_threshold:
             return
 
-        # Create and show the rubber band if it doesn't exist
+        # Create and show the rubber band if not already present
         if not self.rubber_band:
             self.rubber_band = QRubberBand(QRubberBand.Rectangle, self.viewport())
-        
+
         rect = QRect(self.rubber_band_origin, event.pos()).normalized()
         self.rubber_band.setGeometry(rect)
         self.rubber_band.show()
-
-        # Perform dynamic selection on every move
         selection_rect = self.rubber_band.geometry()
         content_widget = self.content_widget
         changed_ids = []
 
+        # Iterate over all annotation widgets to update selection state
         for widget in self.annotation_widgets_by_id.values():
             widget_rect_in_content = widget.geometry()
-            # Map widget's geometry from the content area to the visible viewport
+            # Map widget's rect to viewport coordinates
             widget_rect_in_viewport = QRect(
                 content_widget.mapTo(self.viewport(), widget_rect_in_content.topLeft()),
                 widget_rect_in_content.size()
             )
-
             is_in_band = selection_rect.intersects(widget_rect_in_viewport)
-            
-            # A widget should be selected if it was selected at the start OR is in the band now.
             should_be_selected = (widget in self.selection_at_press) or is_in_band
 
+            # Select or deselect widgets as needed
             if should_be_selected and not widget.is_selected():
                 if self.select_widget(widget):
                     changed_ids.append(widget.data_item.annotation.id)
+
             elif not should_be_selected and widget.is_selected():
                 if self.deselect_widget(widget):
                     changed_ids.append(widget.data_item.annotation.id)
-        
+
+        # Emit signal if any selection state changed
         if changed_ids:
             self.selection_changed.emit(changed_ids)
-        
+
     def mouseReleaseEvent(self, event):
         """Handle mouse release to complete rubber band selection."""
-        # Check if a rubber band drag was in progress
         if self.rubber_band_origin is not None and event.button() == Qt.LeftButton:
             if self.rubber_band and self.rubber_band.isVisible():
                 self.rubber_band.hide()
                 self.rubber_band.deleteLater()
                 self.rubber_band = None
 
-            # **NEEDED CHANGE**: Clean up the stored selection state.
             self.selection_at_press = set()
             self.rubber_band_origin = None
             self.mouse_pressed_on_widget = False
             event.accept()
             return
-            
+
         super().mouseReleaseEvent(event)
 
     def handle_annotation_selection(self, widget, event):
-        """Handle selection of annotation widgets with different modes."""
-        # Get the list of widgets to work with based on isolation mode
-        if self.isolated_mode:
-            # Only work with visible widgets when in isolation mode
-            widget_list = [w for w in self.annotation_widgets_by_id.values() if not w.isHidden()]
-        else:
-            # Use all widgets when not in isolation mode
-            widget_list = list(self.annotation_widgets_by_id.values())
-        
+        """Handle selection of annotation widgets with different modes (single, ctrl, shift)."""
+        widget_list = [w for w in self._get_sorted_widgets() if not w.isHidden()]
+
         try:
             widget_index = widget_list.index(widget)
         except ValueError:
@@ -980,44 +1157,42 @@ class AnnotationViewer(QScrollArea):
         modifiers = event.modifiers()
         changed_ids = []
 
-        # --- The selection logic now identifies which items to change   ---
-        # --- but the core state change happens in select/deselect       ---
-        
+        # Shift or Shift+Ctrl: range selection
         if modifiers == Qt.ShiftModifier or modifiers == (Qt.ShiftModifier | Qt.ControlModifier):
-            # Range selection
             if self.last_selected_index != -1:
-                # Find the last selected widget in the current widget list
+                # Find the last selected widget in the current list
                 last_selected_widget = None
                 for w in self.selected_widgets:
                     if w in widget_list:
                         try:
                             last_index_in_current_list = widget_list.index(w)
-                            if last_selected_widget is None or \
-                               last_index_in_current_list > widget_list.index(last_selected_widget):
+                            if (
+                                last_selected_widget is None
+                                or last_index_in_current_list > widget_list.index(last_selected_widget)
+                            ):
                                 last_selected_widget = w
                         except ValueError:
                             continue
-                
+
                 if last_selected_widget:
                     last_selected_index_in_current_list = widget_list.index(last_selected_widget)
                     start = min(last_selected_index_in_current_list, widget_index)
                     end = max(last_selected_index_in_current_list, widget_index)
                 else:
-                    # Fallback if no previously selected widget is found in current list
-                    start = widget_index
-                    end = widget_index
-                
+                    start, end = widget_index, widget_index
+
+                # Select all widgets in the range
                 for i in range(start, end + 1):
-                    # select_widget will return True if a change occurred
                     if self.select_widget(widget_list[i]):
                         changed_ids.append(widget_list[i].data_item.annotation.id)
             else:
+                # No previous selection, just select the clicked widget
                 if self.select_widget(widget):
                     changed_ids.append(widget.data_item.annotation.id)
                 self.last_selected_index = widget_index
-        
+
+        # Ctrl: toggle selection of the clicked widget
         elif modifiers == Qt.ControlModifier:
-            # Toggle selection
             if widget.is_selected():
                 if self.deselect_widget(widget):
                     changed_ids.append(widget.data_item.annotation.id)
@@ -1025,36 +1200,37 @@ class AnnotationViewer(QScrollArea):
                 if self.select_widget(widget):
                     changed_ids.append(widget.data_item.annotation.id)
             self.last_selected_index = widget_index
-                
+
+        # No modifier: single selection
         else:
-            # Normal click: clear all others and select this one
             newly_selected_id = widget.data_item.annotation.id
-            # Deselect all widgets that are not the clicked one
+
+            # Deselect all others
             for w in list(self.selected_widgets):
                 if w.data_item.annotation.id != newly_selected_id:
                     if self.deselect_widget(w):
                         changed_ids.append(w.data_item.annotation.id)
+
             # Select the clicked widget
             if self.select_widget(widget):
                 changed_ids.append(newly_selected_id)
             self.last_selected_index = widget_index
-            
-        # Update isolation if in isolated mode
+
+        # If in isolated mode, update which widgets are visible
         if self.isolated_mode:
             self._update_isolation()
-        
-        # If any selections were changed, emit the signal
+
+        # Emit signal if any selection state changed
         if changed_ids:
             self.selection_changed.emit(changed_ids)
-            
+
     def _update_isolation(self):
         """Update the isolated view to show only currently selected widgets."""
         if not self.isolated_mode:
             return
-            
+        # If in isolated mode, only show selected widgets
         if self.selected_widgets:
-            # ADD TO isolation instead of replacing it
-            self.isolated_widgets.update(self.selected_widgets)  # Use update() to add, not replace
+            self.isolated_widgets.update(self.selected_widgets)
             self.setUpdatesEnabled(False)
             try:
                 for widget in self.annotation_widgets_by_id.values():
@@ -1063,169 +1239,119 @@ class AnnotationViewer(QScrollArea):
                     else:
                         widget.show()
                 self.recalculate_widget_positions()
+
             finally:
                 self.setUpdatesEnabled(True)
-        else:
-            # If no widgets are selected, keep the current isolation (don't exit)
-            # This prevents accidentally exiting isolation mode when clearing selection
-            pass
 
     def select_widget(self, widget):
-        """Select a widget, update the data_item, and return True if state changed."""
-        if not widget.is_selected():
-            widget.set_selected(True)
+        """Selects a widget, updates its data_item, and returns True if state changed."""
+        if not widget.is_selected():  # is_selected() checks the data_item
+            # 1. Controller modifies the state on the data item
             widget.data_item.set_selected(True)
+            # 2. Controller tells the view to update its appearance
+            widget.update_selection_visuals()
             self.selected_widgets.append(widget)
-            self.update_label_window_selection()
-            self._update_toolbar_state()  # Update button states
+            self._update_toolbar_state()
             return True
         return False
 
     def deselect_widget(self, widget):
-        """Deselect a widget, update the data_item, and return True if state changed."""
+        """Deselects a widget, updates its data_item, and returns True if state changed."""
         if widget.is_selected():
-            widget.set_selected(False)
+            # 1. Controller modifies the state on the data item
             widget.data_item.set_selected(False)
+            # 2. Controller tells the view to update its appearance
+            widget.update_selection_visuals()
             if widget in self.selected_widgets:
                 self.selected_widgets.remove(widget)
-            self.update_label_window_selection()
-            self._update_toolbar_state()  # Update button states
+            self._update_toolbar_state()
             return True
         return False
 
     def clear_selection(self):
         """Clear all selected widgets and update toolbar state."""
         for widget in list(self.selected_widgets):
-            widget.set_selected(False)
-        self.selected_widgets.clear()
-        self.update_label_window_selection()
-        self._update_toolbar_state()  # Update button states
+            # This will internally call deselect_widget, which is fine
+            self.deselect_widget(widget)
 
-    def update_label_window_selection(self):
-        """Update the label window selection based on currently selected annotations."""
-        explorer_window = self.parent()
-        while explorer_window and not hasattr(explorer_window, 'main_window'):
-            explorer_window = explorer_window.parent()
-            
-        if not explorer_window or not hasattr(explorer_window, 'main_window'):
-            return
-            
-        main_window = explorer_window.main_window
-        label_window = main_window.label_window
-        annotation_window = main_window.annotation_window
-        
-        if not self.selected_widgets:
-            label_window.deselect_active_label()
-            label_window.update_annotation_count()
-            return
-            
-        selected_data_items = [widget.data_item for widget in self.selected_widgets]
-        
-        first_effective_label = selected_data_items[0].effective_label
-        all_same_current_label = True
-        for item in selected_data_items:
-            if item.effective_label.id != first_effective_label.id:
-                all_same_current_label = False
-                break
-        
-        if all_same_current_label:
-            label_window.set_active_label(first_effective_label)
-            if not selected_data_items[0].has_preview_changes():
-                annotation_window.labelSelected.emit(first_effective_label.id)
-        else:
-            label_window.deselect_active_label()
-        
-        label_window.update_annotation_count()
+        self.selected_widgets.clear()
+        self._update_toolbar_state()
 
     def get_selected_annotations(self):
         """Get the annotations corresponding to selected widgets."""
         return [widget.annotation for widget in self.selected_widgets]
-    
+
     def render_selection_from_ids(self, selected_ids):
         """Update the visual selection of widgets based on a set of IDs from the controller."""
-        # Block signals temporarily to prevent cascade updates
         self.setUpdatesEnabled(False)
-        
         try:
             for ann_id, widget in self.annotation_widgets_by_id.items():
                 is_selected = ann_id in selected_ids
-                widget.set_selected(is_selected)
-            
-            # Resync internal list of selected widgets
+                # 1. Update the state on the central data item
+                widget.data_item.set_selected(is_selected)
+                # 2. Tell the widget to update its visuals based on the new state
+                widget.update_selection_visuals()
+
+            # Resync internal list of selected widgets from the source of truth
             self.selected_widgets = [w for w in self.annotation_widgets_by_id.values() if w.is_selected()]
-            
-            # If we're in isolated mode, ADD to the isolation instead of replacing it
+
             if self.isolated_mode and self.selected_widgets:
-                self.isolated_widgets.update(self.selected_widgets)  # Add to existing isolation
-                # Hide all widgets except those in the isolated set
+                self.isolated_widgets.update(self.selected_widgets)
                 for widget in self.annotation_widgets_by_id.values():
-                    if widget not in self.isolated_widgets:
-                        widget.hide()
-                    else:
-                        widget.show()
+                    widget.setHidden(widget not in self.isolated_widgets)
                 self.recalculate_widget_positions()
-            
         finally:
             self.setUpdatesEnabled(True)
-        
-        # Update label window once at the end
-        self.update_label_window_selection()
-        # Update toolbar state to enable/disable Isolate button
         self._update_toolbar_state()
-    
+
     def apply_preview_label_to_selected(self, preview_label):
         """Apply a preview label and emit a signal for the embedding view to update."""
         if not self.selected_widgets or not preview_label:
             return
-
         changed_ids = []
         for widget in self.selected_widgets:
             widget.data_item.set_preview_label(preview_label)
-            widget.update() # Force repaint with new color
+            widget.update()  # Force repaint with new color
             changed_ids.append(widget.data_item.annotation.id)
-            
-        # Recalculate positions to update sorting based on new effective labels
+
         if self.sort_combo.currentText() == "Label":
             self.recalculate_widget_positions()
-
         if changed_ids:
             self.preview_changed.emit(changed_ids)
 
     def clear_preview_states(self):
-        """Clear all preview states and revert to original labels."""
-        # We just need to iterate through all widgets and tell their data_items to clear
-        something_cleared = False
+        """
+        Clears all preview states, including label changes,
+        reverting them to their original state.
+        """
+        something_changed = False
         for widget in self.annotation_widgets_by_id.values():
+            # Check for and clear preview labels
             if widget.data_item.has_preview_changes():
                 widget.data_item.clear_preview_label()
-                widget.update() # Repaint to show original color
-                something_cleared = True
-        
-        if something_cleared:
-            # Recalculate positions to update sorting based on reverted labels
-            if self.sort_combo.currentText() == "Label":
+                widget.update()  # Repaint to show original color
+                something_changed = True
+
+        if something_changed:
+            # Recalculate positions to update sorting and re-flow the layout
+            if self.sort_combo.currentText() in ("Label", "Image"):
                 self.recalculate_widget_positions()
-            self.update_label_window_selection()
 
     def has_preview_changes(self):
-        """Check if there are any pending preview changes."""
+        """Return True if there are preview changes."""
         return any(w.data_item.has_preview_changes() for w in self.annotation_widgets_by_id.values())
 
     def get_preview_changes_summary(self):
-        """Get a summary of preview changes for user feedback."""
+        """Get a summary of preview changes."""
         change_count = sum(1 for w in self.annotation_widgets_by_id.values() if w.data_item.has_preview_changes())
-        if not change_count:
-            return "No preview changes"
-        return f"{change_count} annotation(s) with preview changes"
+        return f"{change_count} annotation(s) with preview changes" if change_count else "No preview changes"
 
     def apply_preview_changes_permanently(self):
-        """Apply all preview changes permanently to the annotation data."""
+        """Apply preview changes permanently."""
         applied_annotations = []
         for widget in self.annotation_widgets_by_id.values():
-            # Tell the data_item to apply its changes to the underlying annotation
             if widget.data_item.apply_preview_permanently():
                 applied_annotations.append(widget.annotation)
-        
         return applied_annotations
 
 
@@ -1236,48 +1362,51 @@ class AnnotationViewer(QScrollArea):
 
 class ExplorerWindow(QMainWindow):
     def __init__(self, main_window, parent=None):
+        """Initialize the ExplorerWindow."""
         super(ExplorerWindow, self).__init__(parent)
         self.main_window = main_window
         self.image_window = main_window.image_window
         self.label_window = main_window.label_window
         self.annotation_window = main_window.annotation_window
 
-        self.device = main_window.device  # Use the same device as the main window
-        self.model_path = ""
+        self.device = main_window.device
         self.loaded_model = None
 
-        # Store current filtered data items for embedding
-        self.current_data_items = []
+        self.feature_store = FeatureStore()
         
-        # Cache for extracted features and the model that generated them ---
+        # Add a property to store the parameters with defaults
+        self.mislabel_params = {'k': 20, 'threshold': 0.6}
+        self.uncertainty_params = {'confidence': 0.6, 'margin': 0.1}
+        self.similarity_params = {'k': 30}
+        
+        self.data_item_cache = {}  # Cache for AnnotationDataItem objects
+
+        self.current_data_items = []
         self.current_features = None
         self.current_feature_generating_model = ""
+        self.current_embedding_model_info = None
+        self._ui_initialized = False
 
         self.setWindowTitle("Explorer")
-        # Set the window icon
         explorer_icon_path = get_icon("magic.png")
         self.setWindowIcon(QIcon(explorer_icon_path))
 
-        # Create a central widget and main layout
         self.central_widget = QWidget()
         self.setCentralWidget(self.central_widget)
         self.main_layout = QVBoxLayout(self.central_widget)
-        # Create a left panel widget and layout for the re-parented LabelWindow
         self.left_panel = QWidget()
         self.left_layout = QVBoxLayout(self.left_panel)
 
-        # Create widgets in __init__ so they're always available
-        self.annotation_settings_widget = AnnotationSettingsWidget(self.main_window, self)
-        self.model_settings_widget = ModelSettingsWidget(self.main_window, self)
-        self.embedding_settings_widget = EmbeddingSettingsWidget(self.main_window, self)
-        self.annotation_viewer = AnnotationViewer(self) 
-        self.embedding_viewer = EmbeddingViewer(self)
+        self.annotation_settings_widget = None
+        self.model_settings_widget = None
+        self.embedding_settings_widget = None
+        self.annotation_viewer = None
+        self.embedding_viewer = None
 
-        # Create buttons
         self.clear_preview_button = QPushButton('Clear Preview', self)
         self.clear_preview_button.clicked.connect(self.clear_preview_changes)
         self.clear_preview_button.setToolTip("Clear all preview changes and revert to original labels")
-        self.clear_preview_button.setEnabled(False)  # Initially disabled
+        self.clear_preview_button.setEnabled(False)
 
         self.exit_button = QPushButton('Exit', self)
         self.exit_button.clicked.connect(self.close)
@@ -1286,27 +1415,30 @@ class ExplorerWindow(QMainWindow):
         self.apply_button = QPushButton('Apply', self)
         self.apply_button.clicked.connect(self.apply)
         self.apply_button.setToolTip("Apply changes")
-        self.apply_button.setEnabled(False)  # Initially disabled
+        self.apply_button.setEnabled(False)
 
     def showEvent(self, event):
-        self.setup_ui()
+        """Handle show event."""
+        if not self._ui_initialized:
+            self.setup_ui()
+            self._ui_initialized = True
         super(ExplorerWindow, self).showEvent(event)
 
     def closeEvent(self, event):
-        """
-        Handles the window close event.
-        This now calls the resource cleanup method.
-        """
+        """Handle close event."""
         # Stop any running timers to prevent errors
         if hasattr(self, 'embedding_viewer') and self.embedding_viewer:
             if hasattr(self.embedding_viewer, 'animation_timer') and self.embedding_viewer.animation_timer:
                 self.embedding_viewer.animation_timer.stop()
 
-        # Clear any unsaved preview states
-        if hasattr(self, 'annotation_viewer'):
-            self.annotation_viewer.clear_preview_states()
+        # Call the main cancellation method to revert any pending changes
+        self.clear_preview_changes()
 
-        # --- NEW: Call the dedicated cleanup method ---
+        # Clean up the feature store by deleting its files
+        if hasattr(self, 'feature_store') and self.feature_store:
+            self.feature_store.delete_storage()
+
+        # Call the dedicated cleanup method
         self._cleanup_resources()
 
         # Re-enable the main window before closing
@@ -1319,541 +1451,908 @@ class ExplorerWindow(QMainWindow):
 
         # Clear the reference in the main_window to allow garbage collection
         self.main_window.explorer_window = None
-        
+
+        # Set the ui_initialized flag to False so it can be re-initialized next time
+        self._ui_initialized = False
+
         event.accept()
 
     def setup_ui(self):
-        # Clear the main layout to remove any existing widgets
+        """Set up the UI for the ExplorerWindow."""
         while self.main_layout.count():
             child = self.main_layout.takeAt(0)
             if child.widget():
-                child.widget().setParent(None)  # Remove from layout but don't delete
+                child.widget().setParent(None)
 
-        # Top section: Conditions and Settings side by side
-        top_layout = QHBoxLayout()
+        # 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)
+        if self.annotation_settings_widget is None:
+            self.annotation_settings_widget = AnnotationSettingsWidget(self.main_window, self)
 
-        # Add existing widgets to layout
-        top_layout.addWidget(self.annotation_settings_widget, 2)  # Give annotation settings more space
-        top_layout.addWidget(self.model_settings_widget, 1)      # Model settings in the middle
-        top_layout.addWidget(self.embedding_settings_widget, 1)  # Embedding settings on the right
+        # Model selection panel (choose feature extraction model)
+        if self.model_settings_widget is None:
+            self.model_settings_widget = ModelSettingsWidget(self.main_window, self)
 
-        # Create container widget for top layout
+        # 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)
+
+        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)
 
-        # Middle section: Annotation Viewer (left) and Embedding Viewer (right)
         middle_splitter = QSplitter(Qt.Horizontal)
-
-        # Wrap annotation viewer in a group box
         annotation_group = QGroupBox("Annotation Viewer")
         annotation_layout = QVBoxLayout(annotation_group)
         annotation_layout.addWidget(self.annotation_viewer)
         middle_splitter.addWidget(annotation_group)
 
-        # Wrap embedding viewer in a group box
         embedding_group = QGroupBox("Embedding Viewer")
         embedding_layout = QVBoxLayout(embedding_group)
         embedding_layout.addWidget(self.embedding_viewer)
         middle_splitter.addWidget(embedding_group)
-
-        # Set splitter proportions (annotation viewer wider)
         middle_splitter.setSizes([500, 500])
-
-        # Add middle section to main layout with stretch factor
         self.main_layout.addWidget(middle_splitter, 1)
-
-        # Note: LabelWindow will be re-parented here by MainWindow.open_explorer_window()
-        # The LabelWindow will be added to self.left_layout at index 1 by the MainWindow
         self.main_layout.addWidget(self.label_window)
 
-        # Bottom control buttons
         self.buttons_layout = QHBoxLayout()
-        # Add stretch to push buttons to the right
         self.buttons_layout.addStretch(1)
-
-        # Add existing buttons to layout
         self.buttons_layout.addWidget(self.clear_preview_button)
         self.buttons_layout.addWidget(self.exit_button)
         self.buttons_layout.addWidget(self.apply_button)
-
         self.main_layout.addLayout(self.buttons_layout)
-
-        # Set default condition to current image and refresh filters
+        
+        self._initialize_data_item_cache()
+        self.annotation_settings_widget.set_default_to_current_image()
+        self.refresh_filters()
+        
         self.annotation_settings_widget.set_default_to_current_image()
         self.refresh_filters()
 
-        # Connect label selection to preview updates (only connect once)
         try:
             self.label_window.labelSelected.disconnect(self.on_label_selected_for_preview)
         except TypeError:
-            pass  # Signal wasn't connected yet
+            pass
 
+        # Connect signals to slots
         self.label_window.labelSelected.connect(self.on_label_selected_for_preview)
         self.annotation_viewer.selection_changed.connect(self.on_annotation_view_selection_changed)
         self.annotation_viewer.preview_changed.connect(self.on_preview_changed)
         self.annotation_viewer.reset_view_requested.connect(self.on_reset_view_requested)
         self.embedding_viewer.selection_changed.connect(self.on_embedding_view_selection_changed)
         self.embedding_viewer.reset_view_requested.connect(self.on_reset_view_requested)
-
+        self.embedding_viewer.find_mislabels_requested.connect(self.find_potential_mislabels)
+        self.embedding_viewer.mislabel_parameters_changed.connect(self.on_mislabel_params_changed)
+        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.annotation_viewer.find_similar_requested.connect(self.find_similar_annotations)
+        self.annotation_viewer.similarity_settings_widget.parameters_changed.connect(self.on_similarity_params_changed)
+        
     @pyqtSlot(list)
     def on_annotation_view_selection_changed(self, changed_ann_ids):
-        """A selection was made in the AnnotationViewer, so update the EmbeddingViewer."""
+        """Syncs selection from AnnotationViewer to EmbeddingViewer."""
+        # Per request, unselect any annotation in the main AnnotationWindow
+        if hasattr(self, 'annotation_window'):
+            self.annotation_window.unselect_annotations()
+
         all_selected_ids = {w.data_item.annotation.id for w in self.annotation_viewer.selected_widgets}
-        
-        # Only try to sync the selection with the EmbeddingViewer if it has points.
-        # This prevents the feedback loop that was clearing the selection.
         if self.embedding_viewer.points_by_id:
             self.embedding_viewer.render_selection_from_ids(all_selected_ids)
-            
-        self.update_label_window_selection() # Keep label window in sync
+
+        # Call the new centralized method
+        self.update_label_window_selection()
 
     @pyqtSlot(list)
     def on_embedding_view_selection_changed(self, all_selected_ann_ids):
-        """A selection was made in the EmbeddingViewer, so update the AnnotationViewer."""
-        # Check if this is a new selection being made when nothing was previously selected
+        """Syncs selection from EmbeddingViewer to AnnotationViewer."""
+        # Per request, unselect any annotation in the main AnnotationWindow
+        if hasattr(self, 'annotation_window'):
+            self.annotation_window.unselect_annotations()
+
+        # Check the state BEFORE the selection is changed
         was_empty_selection = len(self.annotation_viewer.selected_widgets) == 0
-        is_new_selection = len(all_selected_ann_ids) > 0
-        
-        # Update the annotation viewer with the new selection
+
+        # Now, update the selection in the annotation viewer
         self.annotation_viewer.render_selection_from_ids(set(all_selected_ann_ids))
-        
-        # Auto-switch to isolation mode if conditions are met
-        if (was_empty_selection and 
-            is_new_selection and 
-            not self.annotation_viewer.isolated_mode):
-            print("Auto-switching to isolation mode due to new selection in embedding viewer")
+
+        # The rest of the logic now works correctly
+        is_new_selection = len(all_selected_ann_ids) > 0
+        if (
+            was_empty_selection and
+            is_new_selection and
+            not self.annotation_viewer.isolated_mode
+        ):
             self.annotation_viewer.isolate_selection()
-        
-        self.update_label_window_selection() # Keep label window in sync
+
+        self.update_label_window_selection()
 
     @pyqtSlot(list)
     def on_preview_changed(self, changed_ann_ids):
-        """A preview color was changed in the AnnotationViewer, so update the EmbeddingViewer points."""
+        """Updates embedding point colors and tooltips when a preview label is applied."""
         for ann_id in changed_ann_ids:
+            # Update embedding point color
             point = self.embedding_viewer.points_by_id.get(ann_id)
             if point:
-                point.update()  # Force the point to repaint itself
-                
+                point.update()
+                point.update_tooltip()  # Refresh tooltip to show new effective label
+
+            # Update annotation widget tooltip
+            widget = self.annotation_viewer.annotation_widgets_by_id.get(ann_id)
+            if widget:
+                widget.update_tooltip()
+
     @pyqtSlot()
     def on_reset_view_requested(self):
         """Handle reset view requests from double-click in either viewer."""
         # Clear all selections in both viewers
         self.annotation_viewer.clear_selection()
         self.embedding_viewer.render_selection_from_ids(set())
-        
-        # Exit isolation mode if currently active
+
+        # Exit isolation mode if currently active in AnnotationViewer
         if self.annotation_viewer.isolated_mode:
             self.annotation_viewer.show_all_annotations()
-        
-        # Update button states
+
+        if self.embedding_viewer.isolated_mode:
+            self.embedding_viewer.show_all_points()
+
+        # Clear similarity sort context
+        self.annotation_viewer.active_ordered_ids = []
+
+        self.update_label_window_selection()
         self.update_button_states()
-        
+
         print("Reset view: cleared selections and exited isolation mode")
+        
+    @pyqtSlot(dict)
+    def on_mislabel_params_changed(self, params):
+        """Updates the stored parameters for mislabel detection."""
+        self.mislabel_params = params
+        print(f"Mislabel detection parameters updated: {self.mislabel_params}")
+        
+    @pyqtSlot(dict)
+    def on_uncertainty_params_changed(self, params):
+        """Updates the stored parameters for uncertainty analysis."""
+        self.uncertainty_params = params
+        print(f"Uncertainty parameters updated: {self.uncertainty_params}")
+        
+    @pyqtSlot(dict)
+    def on_similarity_params_changed(self, params):
+        """Updates the stored parameters for similarity search."""
+        self.similarity_params = params
+        print(f"Similarity search parameters updated: {self.similarity_params}")
+        
+    @pyqtSlot()
+    def on_model_selection_changed(self):
+        """
+        Handles changes in the model settings to enable/disable model-dependent features.
+        """
+        if not self._ui_initialized:
+            return
+
+        model_name, feature_mode = self.model_settings_widget.get_selected_model()
+        is_predict_mode = ".pt" in model_name and feature_mode == "Predictions"
+        
+        self.embedding_viewer.is_uncertainty_analysis_available = is_predict_mode
+        self.embedding_viewer._update_toolbar_state()
+        
+    def _initialize_data_item_cache(self):
+        """
+        Creates a persistent AnnotationDataItem for every annotation,
+        caching them for the duration of the session.
+        """
+        self.data_item_cache.clear()
+        if not hasattr(self.main_window.annotation_window, 'annotations_dict'):
+            return
+
+        all_annotations = self.main_window.annotation_window.annotations_dict.values()
+        for ann in all_annotations:
+            if ann.id not in self.data_item_cache:
+                self.data_item_cache[ann.id] = AnnotationDataItem(ann)
 
     def update_label_window_selection(self):
-        """Update the label window based on the selection in the annotation viewer."""
-        self.annotation_viewer.update_label_window_selection()
+        """
+        Updates the label window based on the selection state of the currently
+        loaded data items. This is the single, centralized point of logic.
+        """
+        # Get selected items directly from the master data list
+        selected_data_items = [
+            item for item in self.current_data_items if item.is_selected
+        ]
+
+        if not selected_data_items:
+            self.label_window.deselect_active_label()
+            self.label_window.update_annotation_count()
+            return
+
+        first_effective_label = selected_data_items[0].effective_label
+        all_same_current_label = all(
+            item.effective_label.id == first_effective_label.id
+            for item in selected_data_items
+        )
+
+        if all_same_current_label:
+            self.label_window.set_active_label(first_effective_label)
+            # This emit is what updates other UI elements, like the annotation list
+            self.annotation_window.labelSelected.emit(first_effective_label.id)
+        else:
+            self.label_window.deselect_active_label()
+
+        self.label_window.update_annotation_count()
 
     def get_filtered_data_items(self):
-        """Get annotations that match all conditions, returned as AnnotationDataItem objects."""
-        data_items = []
+        """
+        Gets annotations matching all conditions by retrieving their
+        persistent AnnotationDataItem objects from the cache.
+        """
         if not hasattr(self.main_window.annotation_window, 'annotations_dict'):
-            return data_items
+            return []
 
-        # Get current filter conditions
         selected_images = self.annotation_settings_widget.get_selected_images()
         selected_types = self.annotation_settings_widget.get_selected_annotation_types()
         selected_labels = self.annotation_settings_widget.get_selected_labels()
 
-        annotations_to_process = []
-        for annotation in self.main_window.annotation_window.annotations_dict.values():
-            annotation_matches = True
-
-            # Check image condition - if empty list, no annotations match
-            if selected_images:
-                annotation_image = os.path.basename(annotation.image_path)
-                if annotation_image not in selected_images:
-                    annotation_matches = False
-            else:
-                # No images selected means no annotations should match
-                annotation_matches = False
-
-            # Check annotation type condition - if empty list, no annotations match
-            if annotation_matches:
-                if selected_types:
-                    annotation_type = type(annotation).__name__
-                    if annotation_type not in selected_types:
-                        annotation_matches = False
-                else:
-                    # No types selected means no annotations should match
-                    annotation_matches = False
-
-            # Check label condition - if empty list, no annotations match
-            if annotation_matches:
-                if selected_labels:
-                    annotation_label = annotation.label.short_label_code
-                    if annotation_label not in selected_labels:
-                        annotation_matches = False
-                else:
-                    # No labels selected means no annotations should match
-                    annotation_matches = False
+        if not all([selected_images, selected_types, selected_labels]):
+            return []
 
-            if annotation_matches:
-                annotations_to_process.append(annotation)
+        annotations_to_process = [
+            ann for ann in self.main_window.annotation_window.annotations_dict.values()
+            if (os.path.basename(ann.image_path) in selected_images and
+                type(ann).__name__ in selected_types and
+                ann.label.short_label_code in selected_labels)
+        ]
 
-        # Ensure all filtered annotations have cropped images
         self._ensure_cropped_images(annotations_to_process)
         
-        # Wrap in AnnotationDataItem
-        for ann in annotations_to_process:
-            data_items.append(AnnotationDataItem(ann))
-
-        return data_items
+        return [self.data_item_cache[ann.id] for ann in annotations_to_process if ann.id in self.data_item_cache]
     
-    def _ensure_cropped_images(self, annotations):
-        """Ensure all provided annotations have a cropped image available."""
-        annotations_by_image = {}
-        for annotation in annotations:
-            # Only process annotations that don't have a cropped image yet
-            if not annotation.cropped_image:
-                image_path = annotation.image_path
-                if image_path not in annotations_by_image:
-                    annotations_by_image[image_path] = []
-                annotations_by_image[image_path].append(annotation)
-        
-        # Only proceed if there are annotations that actually need cropping
-        if annotations_by_image:
-            progress_bar = ProgressBar(self, "Cropping Image Annotations")
-            progress_bar.show()
-            progress_bar.start_progress(len(annotations_by_image))
-            
-            try:
-                # Crop annotations for each image using the AnnotationWindow method
-                # This ensures consistency with how cropped images are generated elsewhere
-                for image_path, image_annotations in annotations_by_image.items():
-                    self.annotation_window.crop_annotations(
-                        image_path=image_path, 
-                        annotations=image_annotations, 
-                        return_annotations=False,  # We don't need the return value
-                        verbose=False
-                    )
-                    # Update progress bar
-                    progress_bar.update_progress()
-                    
-            except Exception as e:
-                print(f"Error cropping annotations: {e}")
-                
-            finally:
-                progress_bar.finish_progress()
-                progress_bar.stop_progress()
-                progress_bar.close()
-
-    def _extract_color_features(self, data_items, progress_bar=None, bins=32):
+    def find_potential_mislabels(self):
         """
-        Extracts a comprehensive set of color features from annotation crops using only NumPy.
-        
-        For each image, it calculates:
-        1.  Color Moments (per channel):
-            - Mean (1st moment)
-            - Standard Deviation (2nd moment)
-            - Skewness (3rd moment)
-            - Kurtosis (4th moment)
-        2.  Color Histogram (per channel)
-        3.  Grayscale Statistics:
-            - Mean Brightness
-            - Contrast (Std Dev)
-            - Intensity Range
-        4.  Geometric Features:
-            - Area
-            - Perimeter
+        Identifies annotations whose label does not match the majority of its
+        k-nearest neighbors in the high-dimensional feature space.
         """
-        if progress_bar:
-            progress_bar.set_title("Extracting Color Features...")
-            progress_bar.start_progress(len(data_items))
+        # Get parameters from the stored property instead of hardcoding
+        K = self.mislabel_params.get('k', 5)
+        agreement_threshold = self.mislabel_params.get('threshold', 0.6)
 
-        features = []
-        valid_data_items = []
-        for item in data_items:
-            pixmap = item.annotation.get_cropped_image()
-            if pixmap and not pixmap.isNull():
-                # arr has shape (height, width, 3)
-                arr = pixmap_to_numpy(pixmap)
-                
-                # Reshape for channel-wise statistics: (num_pixels, 3)
-                pixels = arr.reshape(-1, 3)
+        if not self.embedding_viewer.points_by_id or len(self.embedding_viewer.points_by_id) < K:
+            QMessageBox.information(self, "Not Enough Data",
+                                    f"This feature requires at least {K} points in the embedding viewer.")
+            return
 
-                # --- 1. Calculate Color Moments using only NumPy ---
-                mean_color = np.mean(pixels, axis=0)
-                std_color = np.std(pixels, axis=0)
+        items_in_view = list(self.embedding_viewer.points_by_id.values())
+        data_items_in_view = [p.data_item for p in items_in_view]
+
+        # Get the model key used for the current embedding
+        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(' ', '_')
+        # FIX: Also replace the forward slash to handle "N/A"
+        sanitized_feature_mode = feature_mode.replace(' ', '_').replace('/', '_')
+        model_key = f"{sanitized_model_name}_{sanitized_feature_mode}"
+
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        try:
+            # Get the FAISS index and the mapping from index to annotation ID
+            index = self.feature_store._get_or_load_index(model_key)
+            faiss_idx_to_ann_id = self.feature_store.get_faiss_index_to_annotation_id_map(model_key)
+            if index is None or not faiss_idx_to_ann_id:
+                QMessageBox.warning(self, "Error", "Could not find a valid feature index for the current model.")
+                return
+
+            # Get the high-dimensional features for the points in the current view
+            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')
+
+            # Perform k-NN search. We search for K+1 because the point itself will be the first result.
+            _, I = index.search(query_vectors, K + 1)
+
+            mislabeled_ann_ids = []
+            for i, ann_id in enumerate(query_ann_ids):
+                current_label = self.data_item_cache[ann_id].effective_label.id
                 
-                # Center the data (subtract the mean) for skew/kurtosis calculation
-                centered_pixels = pixels - mean_color
+                # Get neighbor labels, ignoring the first result (the point itself)
+                neighbor_faiss_indices = I[i][1:]
                 
-                # A small value (epsilon) is added to the denominator to prevent division by zero
-                epsilon = 1e-8
-                skew_color = np.mean(centered_pixels**3, axis=0) / (std_color**3 + epsilon)
-                kurt_color = np.mean(centered_pixels**4, axis=0) / (std_color**4 + epsilon) - 3
+                neighbor_labels = []
+                for n_idx in neighbor_faiss_indices:
+                    # THIS IS THE CORRECTED LOGIC
+                    if n_idx in faiss_idx_to_ann_id:
+                        neighbor_ann_id = faiss_idx_to_ann_id[n_idx]
+                        # ADD THIS CHECK to ensure the neighbor hasn't been deleted
+                        if neighbor_ann_id in self.data_item_cache:
+                            neighbor_labels.append(self.data_item_cache[neighbor_ann_id].effective_label.id)
+
+                if not neighbor_labels:
+                    continue
+
+                # Use the agreement threshold instead of strict majority
+                num_matching_neighbors = neighbor_labels.count(current_label)
+                agreement_ratio = num_matching_neighbors / len(neighbor_labels)
+
+                if agreement_ratio < agreement_threshold:
+                    mislabeled_ann_ids.append(ann_id)
+
+            self.embedding_viewer.render_selection_from_ids(set(mislabeled_ann_ids))
+
+        finally:
+            QApplication.restoreOverrideCursor()
+            
+    def find_uncertain_annotations(self):
+        """
+        Identifies annotations where the model's prediction is uncertain.
+        It reuses cached predictions if available, otherwise runs a temporary prediction.
+        """
+        if not self.embedding_viewer.points_by_id:
+            QMessageBox.information(self, "No Data", "Please generate an embedding first.")
+            return
+
+        if self.current_embedding_model_info is None:
+            QMessageBox.information(self, 
+                                    "No Embedding", 
+                                    "Could not determine the model used for the embedding. Please run it again.")
+            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_name_from_embedding, feature_mode_from_embedding = self.current_embedding_model_info
+        
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        try:
+            probabilities_dict = {}
+
+            # Decide whether to reuse cached features or run a new prediction
+            if feature_mode_from_embedding == "Predictions":
+                print("Reusing cached prediction vectors from the FeatureStore.")
+                sanitized_model_name = os.path.basename(model_name_from_embedding).replace(' ', '_').replace('/', '_')
+                sanitized_feature_mode = feature_mode_from_embedding.replace(' ', '_').replace('/', '_')
+                model_key = f"{sanitized_model_name}_{sanitized_feature_mode}"
                 
-                # --- 2. Calculate Color Histograms ---
-                histograms = []
-                for i in range(3): # For each channel (R, G, B)
-                    hist, _ = np.histogram(pixels[:, i], bins=bins, range=(0, 255))
-                    
-                    # Normalize histogram
-                    hist_sum = np.sum(hist)
-                    if hist_sum > 0:
-                        histograms.append(hist / hist_sum)
-                    else: # Avoid division by zero
-                        histograms.append(np.zeros(bins))
+                probabilities_dict, _ = self.feature_store.get_features(data_items_in_view, model_key)
+                if not probabilities_dict:
+                    QMessageBox.warning(self, 
+                                        "Cache Error", 
+                                        "Could not retrieve cached predictions.")
+                    return
+            else:
+                print("Embedding not based on 'Predictions' mode. Running a temporary prediction.")
+                model_info_for_predict = self.model_settings_widget.get_selected_model()
+                probabilities_dict = self._get_yolo_predictions_for_uncertainty(data_items_in_view, 
+                                                                                model_info_for_predict)
+
+            if not probabilities_dict:
+                # The helper function will show its own, more specific errors.
+                return
+
+            uncertain_ids = []
+            params = self.uncertainty_params
+            for ann_id, probs in probabilities_dict.items():
+                if len(probs) < 2:
+                    continue  # Cannot calculate margin
+
+                sorted_probs = np.sort(probs)[::-1]
+                top1_conf = sorted_probs[0]
+                top2_conf = sorted_probs[1]
+                margin = top1_conf - top2_conf
+
+                if top1_conf < params['confidence'] or margin < params['margin']:
+                    uncertain_ids.append(ann_id)
+            
+            self.embedding_viewer.render_selection_from_ids(set(uncertain_ids))
+            print(f"Found {len(uncertain_ids)} uncertain annotations.")
+
+        finally:
+            QApplication.restoreOverrideCursor()
+            
+    @pyqtSlot()
+    def find_similar_annotations(self):
+        """
+        Finds k-nearest neighbors to the selected annotation(s) and updates 
+        the UI to show the results in an isolated, ordered view. This method
+        now ensures the grid is always updated and resets the sort-by dropdown.
+        """
+        k = self.similarity_params.get('k', 10)
+
+        if not self.annotation_viewer.selected_widgets:
+            QMessageBox.information(self, "No Selection", "Please select one or more annotations first.")
+            return
+
+        if not self.current_embedding_model_info:
+            QMessageBox.warning(self, "No Embedding", "Please run an embedding before searching for similar items.")
+            return
+
+        selected_data_items = [widget.data_item for widget in self.annotation_viewer.selected_widgets]
+        model_name, feature_mode = self.current_embedding_model_info
+        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}"
+
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        try:
+            features_dict, _ = self.feature_store.get_features(selected_data_items, model_key)
+            if not features_dict:
+                QMessageBox.warning(self, 
+                                    "Features Not Found", 
+                                    "Could not retrieve feature vectors for the selected items.")
+                return
+
+            source_vectors = np.array(list(features_dict.values()))
+            query_vector = np.mean(source_vectors, axis=0, keepdims=True).astype('float32')
+
+            index = self.feature_store._get_or_load_index(model_key)
+            faiss_idx_to_ann_id = self.feature_store.get_faiss_index_to_annotation_id_map(model_key)
+            if index is None or not faiss_idx_to_ann_id:
+                QMessageBox.warning(self, 
+                                    "Index Error", 
+                                    "Could not find a valid feature index for the current model.")
+                return
+
+            # Find k results, plus more to account for the query items possibly being in the results
+            num_to_find = k + len(selected_data_items)
+            if num_to_find > index.ntotal:
+                num_to_find = index.ntotal
+            
+            _, I = index.search(query_vector, num_to_find)
+
+            source_ids = {item.annotation.id for item in selected_data_items}
+            similar_ann_ids = []
+            for faiss_idx in I[0]:
+                ann_id = faiss_idx_to_ann_id.get(faiss_idx)
+                if ann_id and ann_id in self.data_item_cache and ann_id not in source_ids:
+                    similar_ann_ids.append(ann_id)
+                if len(similar_ann_ids) == k:
+                    break
+
+            # Create the final ordered list: original selection first, then similar items.
+            ordered_ids_to_display = list(source_ids) + similar_ann_ids
+            
+            # --- FIX IMPLEMENTATION ---
+            # 1. Force sort combo to "None" to avoid user confusion.
+            self.annotation_viewer.sort_combo.setCurrentText("None")
+
+            # 2. Update the embedding viewer selection.
+            self.embedding_viewer.render_selection_from_ids(set(ordered_ids_to_display))
+            
+            # 3. Call the new robust method in AnnotationViewer to handle isolation and grid updates.
+            self.annotation_viewer.display_and_isolate_ordered_results(ordered_ids_to_display)
+
+            self.update_button_states()
+
+        finally:
+            QApplication.restoreOverrideCursor()
+            
+    def _get_yolo_predictions_for_uncertainty(self, data_items, model_info):
+        """
+        Runs a YOLO classification model to get probabilities for uncertainty analysis.
+        This is a streamlined method that does NOT use the feature store.
+        """
+        model_name, feature_mode = model_info
+        
+        # Load the model
+        model, imgsz = self._load_yolo_model(model_name, feature_mode)
+        if model is None:
+            QMessageBox.warning(self, 
+                                "Model Load Error",
+                                f"Could not load YOLO model '{model_name}'.")
+            return None
+        
+        # Prepare images from data items
+        image_list, valid_data_items = self._prepare_images_from_data_items(data_items)
+        if not image_list:
+            return None
+        
+        try:
+            # We need probabilities for uncertainty analysis, so we always use predict
+            results = model.predict(image_list, 
+                                    stream=False,  # Use batch processing for uncertainty
+                                    imgsz=imgsz, 
+                                    half=True, 
+                                    device=self.device, 
+                                    verbose=False)
                 
-                # --- 3. Calculate Grayscale Statistics ---
-                gray_arr = np.dot(arr[..., :3], [0.2989, 0.5870, 0.1140])
-                grayscale_stats = np.array([
-                    np.mean(gray_arr),  # Overall brightness
-                    np.std(gray_arr),   # Overall contrast
-                    np.ptp(gray_arr)    # Peak-to-peak intensity range
-                ])
+            _, probabilities_dict = self._process_model_results(results, valid_data_items, "Predictions")
+            return probabilities_dict
+            
+        except TypeError:
+            QMessageBox.warning(self, 
+                                "Invalid Model",
+                                "The selected model is not compatible with uncertainty analysis.")
+            return None
+            
+        finally:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def _ensure_cropped_images(self, annotations):
+        """Ensures all provided annotations have a cropped image available."""
+        annotations_by_image = {}
+
+        for annotation in annotations:
+            if not annotation.cropped_image:
+                image_path = annotation.image_path
+                if image_path not in annotations_by_image:
+                    annotations_by_image[image_path] = []
+                annotations_by_image[image_path].append(annotation)
+
+        if not annotations_by_image:
+            return
+
+        progress_bar = ProgressBar(self, "Cropping Image Annotations")
+        progress_bar.show()
+        progress_bar.start_progress(len(annotations_by_image))
+
+        try:
+            for image_path, image_annotations in annotations_by_image.items():
+                self.annotation_window.crop_annotations(image_path=image_path,
+                                                        annotations=image_annotations,
+                                                        return_annotations=False,
+                                                        verbose=False)
+                progress_bar.update_progress()
+        finally:
+            progress_bar.finish_progress()
+            progress_bar.stop_progress()
+            progress_bar.close()
+            
+    def _load_yolo_model(self, model_name, feature_mode):
+        """
+        Helper function to load a YOLO model and cache it.
+        
+        Args:
+            model_name (str): Path to the YOLO model file
+            feature_mode (str): Mode for feature extraction ("Embed Features" or "Predictions")
+        
+        Returns:
+            tuple: (model, image_size) or (None, None) if loading fails
+        """
+        current_run_key = (model_name, feature_mode)
+        
+        # Force a reload if the model path OR the feature mode has changed
+        if current_run_key != self.current_feature_generating_model or self.loaded_model is None:
+            print(f"Model or mode changed. Reloading {model_name} for '{feature_mode}'.")
+            try:
+                model = YOLO(model_name)
+                # Update the cache key to the new successful combination
+                self.current_feature_generating_model = current_run_key
+                self.loaded_model = model
+                imgsz = getattr(model.model.args, 'imgsz', 128)
                 
-                # --- 4. Calculate Geometric Features ---
-                area = getattr(item.annotation, 'area', 0.0)
-                perimeter = getattr(item.annotation, 'perimeter', 0.0)
-                geometric_features = np.array([area, perimeter])
+                # Warm up the model
+                dummy_image = np.zeros((imgsz, imgsz, 3), dtype=np.uint8)
+                model.predict(dummy_image, imgsz=imgsz, half=True, device=self.device, verbose=False)
                 
-                # --- 5. Concatenate all features into a single vector ---
-                current_features = np.concatenate([
-                    mean_color,
-                    std_color,
-                    skew_color,
-                    kurt_color,
-                    *histograms,
-                    grayscale_stats,
-                    geometric_features
-                ])
+                return model, imgsz
                 
-                features.append(current_features)
+            except Exception as e:
+                print(f"ERROR: Could not load YOLO model '{model_name}': {e}")
+                # On failure, reset the model cache
+                self.loaded_model = None
+                self.current_feature_generating_model = None
+                return None, None
+        
+        # Model already loaded and cached
+        return self.loaded_model, getattr(self.loaded_model.model.args, 'imgsz', 128)
+
+    def _prepare_images_from_data_items(self, data_items, progress_bar=None):
+        """
+        Prepare images from data items for model prediction.
+        
+        Args:
+            data_items (list): List of AnnotationDataItem objects
+            progress_bar (ProgressBar, optional): Progress bar for UI updates
+        
+        Returns:
+            tuple: (image_list, valid_data_items)
+        """
+        if progress_bar:
+            progress_bar.set_title("Preparing images...")
+            progress_bar.start_progress(len(data_items))
+        
+        image_list, valid_data_items = [], []
+        for item in data_items:
+            pixmap = item.annotation.get_cropped_image()
+            if pixmap and not pixmap.isNull():
+                image_list.append(pixmap_to_numpy(pixmap))
                 valid_data_items.append(item)
-            else:
-                print(f"Warning: Could not get cropped image for annotation ID {item.annotation.id}. Skipping.")
             
             if progress_bar:
                 progress_bar.update_progress()
+        
+        return image_list, valid_data_items
 
-        return np.array(features), valid_data_items
-
-    def _extract_yolo_features(self, data_items, model_info, progress_bar=None):
+    def _process_model_results(self, results, valid_data_items, feature_mode, progress_bar=None):
         """
-        Extracts features from annotation crops using a specified YOLO model.
-        Uses model.embed() for embedding features or model.predict() for classification probabilities.
+        Process model results and update data item tooltips.
+        
+        Args:
+            results: Model prediction results
+            valid_data_items (list): List of valid data items
+            feature_mode (str): Mode for feature extraction
+            progress_bar (ProgressBar, optional): Progress bar for UI updates
+        
+        Returns:
+            tuple: (features_list, probabilities_dict)
         """
-        # Unpack model information
-        model_name, feature_mode = model_info
+        features_list = []
+        probabilities_dict = {}
         
-        # Load or retrieve the cached model
-        if model_name != self.model_path or self.loaded_model is None:
-            try:
-                print(f"Loading new model: {model_name}")
-                self.loaded_model = YOLO(model_name)
-                self.model_path = model_name
+        # Get class names from the model for better tooltips
+        model = self.loaded_model.model if hasattr(self.loaded_model, 'model') else None
+        class_names = model.names if model and hasattr(model, 'names') else {}
+        
+        for i, result in enumerate(results):
+            if i >= len(valid_data_items):
+                break
                 
-                # Determine image size from model config if possible
-                try:
-                    self.imgsz = self.loaded_model.model.args['imgsz']
-                    if self.imgsz > 224:
-                        self.imgsz = 128
-                except (AttributeError, KeyError):
-                    self.imgsz = 128
+            item = valid_data_items[i]
+            ann_id = item.annotation.id
+            
+            if feature_mode == "Embed Features":
+                embedding = result.cpu().numpy().flatten()
+                features_list.append(embedding)
+                
+            elif hasattr(result, 'probs') and result.probs is not None:
+                probs = result.probs.data.cpu().numpy().squeeze()
+                features_list.append(probs)
+                probabilities_dict[ann_id] = probs
+                
+                # Store the probabilities directly on the data item for confidence sorting
+                item.prediction_probabilities = probs
                 
-                # Run a dummy inference to warm up the model
-                print(f"Warming up model on device '{self.device}'...")
-                dummy_image = np.zeros((self.imgsz, self.imgsz, 3), dtype=np.uint8)
-                self.loaded_model.predict(dummy_image, imgsz=self.imgsz, half=True, device=self.device, verbose=False)
+                # Format and store prediction details for tooltips
+                if len(probs) > 0:
+                    # Get top 5 predictions
+                    top_indices = probs.argsort()[::-1][:5]
+                    top_probs = probs[top_indices]
                     
-            except Exception as e:
-                print(f"ERROR: Could not load YOLO model '{model_name}': {e}")
-                return np.array([]), []
+                    formatted_preds = ["<b>Top Predictions:</b>"]
+                    for idx, prob in zip(top_indices, top_probs):
+                        class_name = class_names.get(int(idx), f"Class {idx}")
+                        formatted_preds.append(f"{class_name}: {prob*100:.1f}%")
+                    
+                    item.prediction_details = "<br>".join(formatted_preds)
+            else:
+                raise TypeError(
+                    "The 'Predictions' feature mode requires a classification model "
+                    "(e.g., 'yolov8n-cls.pt') that returns class probabilities. "
+                    "The selected model did not provide this output. "
+                    "Please use 'Embed Features' mode for this model."
+                )
                 
+            if progress_bar:
+                progress_bar.update_progress()
+        
+        # After processing is complete, update tooltips
+        for item in valid_data_items:
+            if hasattr(item, 'update_tooltip'):
+                item.update_tooltip()
+                
+        return features_list, probabilities_dict
+
+    def _extract_color_features(self, data_items, progress_bar=None, bins=32):
+        """
+        Extracts color-based features from annotation crops.
+
+        Features extracted per annotation:
+            - Mean, standard deviation, skewness, and kurtosis for each RGB channel
+            - Normalized histogram for each RGB channel
+            - Grayscale statistics: mean, std, range
+            - Geometric features: area, perimeter (if available)
+        Returns:
+            features: np.ndarray of shape (N, feature_dim)
+            valid_data_items: list of AnnotationDataItem with valid crops
+        """
         if progress_bar:
-            progress_bar.set_title(f"Preparing images...")
+            progress_bar.set_title("Extracting features...")
             progress_bar.start_progress(len(data_items))
 
-        # 1. Prepare a list of all valid images and their corresponding data items.
-        image_list = []
+        features = []
         valid_data_items = []
+
         for item in data_items:
             pixmap = item.annotation.get_cropped_image()
             if pixmap and not pixmap.isNull():
-                image_np = pixmap_to_numpy(pixmap)
-                image_list.append(image_np)
+                # Convert QPixmap to numpy array (H, W, 3)
+                arr = pixmap_to_numpy(pixmap)
+                pixels = arr.reshape(-1, 3)
+
+                # Basic color statistics
+                mean_color = np.mean(pixels, axis=0)
+                std_color = np.std(pixels, axis=0)
+
+                # Skewness and kurtosis for each channel
+                epsilon = 1e-8  # Prevent division by zero
+                centered_pixels = pixels - mean_color
+                skew_color = np.mean(centered_pixels ** 3, axis=0) / (std_color ** 3 + epsilon)
+                kurt_color = np.mean(centered_pixels ** 4, axis=0) / (std_color ** 4 + epsilon) - 3
+
+                # Normalized histograms for each channel
+                histograms = [
+                    np.histogram(pixels[:, i], bins=bins, range=(0, 255))[0]
+                    for i in range(3)
+                ]
+                histograms = [
+                    h / h.sum() if h.sum() > 0 else np.zeros(bins)
+                    for h in histograms
+                ]
+
+                # Grayscale statistics
+                gray_arr = np.dot(arr[..., :3], [0.2989, 0.5870, 0.1140])
+                grayscale_stats = np.array([
+                    np.mean(gray_arr),
+                    np.std(gray_arr),
+                    np.ptp(gray_arr)
+                ])
+
+                # Geometric features (area, perimeter)
+                area = getattr(item.annotation, 'area', 0.0)
+                perimeter = getattr(item.annotation, 'perimeter', 0.0)
+                geometric_features = np.array([area, perimeter])
+
+                # Concatenate all features into a single vector
+                current_features = np.concatenate([
+                    mean_color,
+                    std_color,
+                    skew_color,
+                    kurt_color,
+                    *histograms,
+                    grayscale_stats,
+                    geometric_features
+                ])
+
+                features.append(current_features)
                 valid_data_items.append(item)
-            else:
-                print(f"Warning: Could not get cropped image for annotation ID {item.annotation.id}. Skipping.")
 
             if progress_bar:
                 progress_bar.update_progress()
 
+        return np.array(features), valid_data_items
+
+    def _extract_yolo_features(self, data_items, model_info, progress_bar=None):
+        """Extracts features from annotation crops using a YOLO model."""
+        model_name, feature_mode = model_info
+        
+        # Load the model
+        model, imgsz = self._load_yolo_model(model_name, feature_mode)
+        if model is None:
+            return np.array([]), []
+        
+        # Prepare images from data items
+        image_list, valid_data_items = self._prepare_images_from_data_items(data_items, progress_bar)
         if not valid_data_items:
-            print("Warning: No valid images found to process.")
             return np.array([]), []
-
-        embeddings_list = []
+        
+        # Set up prediction parameters
+        kwargs = {
+            'stream': True,
+            'imgsz': imgsz,
+            'half': True,
+            'device': self.device,
+            'verbose': False
+        }
+        
+        # Get results based on feature mode
+        if feature_mode == "Embed Features":
+            results_generator = model.embed(image_list, **kwargs)
+        else:
+            results_generator = model.predict(image_list, **kwargs)
+        
+        if progress_bar:
+            progress_bar.set_title("Extracting features...")
+            progress_bar.start_progress(len(valid_data_items))
         
         try:
-            if progress_bar:
-                progress_bar.set_busy_mode(f"Extracting features with {os.path.basename(model_name)}...")
-            
-                kwargs = {
-                    'stream': True,
-                    'imgsz': self.imgsz,
-                    'half': True,
-                    'device': self.device,
-                    'verbose': False
-                }
-                            
-            # 2. Choose between embed() and predict() based on feature mode
-            using_embed_method = feature_mode == "Embed Features"
+            features_list, _ = self._process_model_results(results_generator, 
+                                                           valid_data_items, 
+                                                           feature_mode,
+                                                           progress_bar=progress_bar)
             
-            if using_embed_method:
-                print("Using embed() method")
-                
-                # Use model.embed() method (uses the second to last layer)
-                results_generator = self.loaded_model.embed(image_list, **kwargs)
-                
-            else:
-                print("Using predict() method")
-                
-                # Use model.predict() method for classification probabilities
-                results_generator = self.loaded_model.predict(image_list, **kwargs)
-            
-            if progress_bar:
-                progress_bar.set_title(f"Extracting features with {os.path.basename(model_name)}...")
-                progress_bar.start_progress(len(valid_data_items))
+            return np.array(features_list), valid_data_items
             
-            # 3. Process the results from the generator - different handling based on method
-            for result in results_generator:
-                if using_embed_method:
-                    try:
-                        # With embed(), result is directly a tensor
-                        embedding = result.cpu().numpy().flatten()
-                        embeddings_list.append(embedding)
-                    except Exception as e:
-                        print(f"Error processing embedding: {e}")
-                        raise TypeError(
-                            f"Model '{os.path.basename(model_name)}' did not return valid embeddings. "
-                            f"Error: {str(e)}"
-                        )
-                else:
-                    # Classification mode: We expect probability vectors
-                    if hasattr(result, 'probs') and result.probs is not None:
-                        embedding = result.probs.data.cpu().numpy().squeeze()
-                        embeddings_list.append(embedding)
-                    else:
-                        raise TypeError(
-                            f"Model '{os.path.basename(model_name)}' did not return probability vectors. "
-                            "Make sure this is a classification model."
-                        )
-
-                # Update progress for each item as it's processed from the stream.
-                if progress_bar:
-                    progress_bar.update_progress()
-
-            if not embeddings_list:
-                print("Warning: No features were extracted. The model may have failed.")
-                return np.array([]), []
-
-            embeddings = np.array(embeddings_list)
-        except Exception as e:
-            print(f"ERROR: An error occurred during feature extraction: {e}")
-            return np.array([]), []
         finally:
-            # Clean up CUDA memory after the operation
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
-        
-        print(f"Successfully extracted {len(embeddings)} features with shape {embeddings.shape}")
-        return embeddings, valid_data_items
 
     def _extract_features(self, data_items, progress_bar=None):
-        """
-        Dispatcher method to call the appropriate feature extraction function.
-        It now passes the progress_bar object to the sub-methods.
-        """
+        """Dispatcher to call the appropriate feature extraction function."""
+        # Get the selected model and feature mode from the model settings widget
         model_name, feature_mode = self.model_settings_widget.get_selected_model()
-        
-        # Handle tuple or string return value
-        if isinstance(model_name, tuple) and len(model_name) >= 3:
+
+        if isinstance(model_name, tuple):
             model_name = model_name[0]
-        else:
-            model_name = model_name
 
         if not model_name:
-            print("No model selected or path provided.")
             return np.array([]), []
 
-        # --- MODIFIED: Pass the progress_bar object ---
         if model_name == "Color Features":
             return self._extract_color_features(data_items, progress_bar=progress_bar)
+
         elif ".pt" in model_name:
-            # Pass the full model_info which may include embed layers
             return self._extract_yolo_features(data_items, (model_name, feature_mode), progress_bar=progress_bar)
-        else:
-            print(f"Unknown or invalid feature model selected: {model_name}")
-            return np.array([]), []
+
+        return np.array([]), []
 
     def _run_dimensionality_reduction(self, features, params):
         """
-        Runs PCA, UMAP or t-SNE on the feature matrix using provided parameters.
+        Runs dimensionality reduction with automatic PCA preprocessing for UMAP and t-SNE.
+
+        Args:
+            features (np.ndarray): Feature matrix of shape (N, D).
+            params (dict): Embedding parameters, including technique and its hyperparameters.
+
+        Returns:
+            np.ndarray or None: 2D embedded features of shape (N, 2), or None on failure.
         """
-        technique = params.get('technique', 'UMAP')  # Changed default to UMAP
-        random_state = 42
-        
-        print(f"Running {technique} on {len(features)} items with params: {params}")
-        if len(features) <= 2:  # UMAP/t-SNE need at least a few points
-            print("Not enough data points for dimensionality reduction.")
+        technique = params.get('technique', 'UMAP')
+        # Default number of components to use for PCA preprocessing
+        pca_components = params.get('pca_components', 50)
+
+        if len(features) <= 2:
+            # Not enough samples for dimensionality reduction
             return None
 
         try:
-            # Scaling is crucial, your implementation is already correct
-            scaler = StandardScaler()
-            features_scaled = scaler.fit_transform(features)
-            
+            # Standardize features before reduction
+            features_scaled = StandardScaler().fit_transform(features)
+            
+            # Apply PCA preprocessing automatically for UMAP or TSNE
+            # (only if the feature dimension is larger than the target PCA components)
+            if technique in ["UMAP", "TSNE"] and features_scaled.shape[1] > pca_components:
+                # Ensure pca_components doesn't exceed number of samples or features
+                pca_components = min(pca_components, features_scaled.shape[0] - 1, features_scaled.shape[1])
+                print(f"Applying PCA preprocessing to {pca_components} components before {technique}")
+                pca = PCA(n_components=pca_components, random_state=42)
+                features_scaled = pca.fit_transform(features_scaled)
+                variance_explained = sum(pca.explained_variance_ratio_) * 100
+                print(f"Variance explained by PCA: {variance_explained:.1f}%")
+
+            # Proceed with the selected dimensionality reduction technique
             if technique == "UMAP":
-                # Get hyperparameters from params, with sensible defaults
-                n_neighbors = params.get('n_neighbors', 15)
-                min_dist = params.get('min_dist', 0.1)
-                metric = params.get('metric', 'cosine')  # Allow metric to be specified
-
+                n_neighbors = min(params.get('n_neighbors', 15), len(features_scaled) - 1)
                 reducer = UMAP(
-                    n_components=2, 
-                    random_state=random_state, 
-                    n_neighbors=min(n_neighbors, len(features_scaled) - 1),
-                    min_dist=min_dist,
-                    metric=metric  # Use the specified metric
+                    n_components=2,
+                    random_state=42,
+                    n_neighbors=n_neighbors,
+                    min_dist=params.get('min_dist', 0.1),
+                    metric=params.get('metric', 'cosine')
                 )
-            
             elif technique == "TSNE":
-                perplexity = params.get('perplexity', 30)
-                early_exaggeration = params.get('early_exaggeration', 12.0)
-                learning_rate = params.get('learning_rate', 'auto')
-
+                perplexity = min(params.get('perplexity', 30), len(features_scaled) - 1)
                 reducer = TSNE(
-                    n_components=2, 
-                    random_state=random_state, 
-                    perplexity=min(perplexity, len(features_scaled) - 1),
-                    early_exaggeration=early_exaggeration,
-                    learning_rate=learning_rate,
-                    init='pca'  # Improves stability and speed
+                    n_components=2,
+                    random_state=42,
+                    perplexity=perplexity,
+                    early_exaggeration=params.get('early_exaggeration', 12.0),
+                    learning_rate=params.get('learning_rate', 'auto'),
+                    init='pca'
                 )
-                
             elif technique == "PCA":
-                reducer = PCA(n_components=2, random_state=random_state)
-
+                reducer = PCA(n_components=2, random_state=42)
             else:
-                print(f"Unknown dimensionality reduction technique: {technique}")
                 return None
-                
+
+            # Fit and transform the features
             return reducer.fit_transform(features_scaled)
-            
+
         except Exception as e:
             print(f"Error during {technique} dimensionality reduction: {e}")
             return None
@@ -1861,91 +2360,115 @@ class ExplorerWindow(QMainWindow):
     def _update_data_items_with_embedding(self, data_items, embedded_features):
         """Updates AnnotationDataItem objects with embedding results."""
         scale_factor = 4000
-        min_vals = np.min(embedded_features, axis=0)
-        max_vals = np.max(embedded_features, axis=0)
+        min_vals, max_vals = np.min(embedded_features, axis=0), np.max(embedded_features, axis=0)
         range_vals = max_vals - min_vals
-
         for i, item in enumerate(data_items):
-            # Normalize coordinates for consistent display
             norm_x = (embedded_features[i, 0] - min_vals[0]) / range_vals[0] if range_vals[0] > 0 else 0.5
             norm_y = (embedded_features[i, 1] - min_vals[1]) / range_vals[1] if range_vals[1] > 0 else 0.5
-            # Scale and center the points in the view
             item.embedding_x = (norm_x * scale_factor) - (scale_factor / 2)
             item.embedding_y = (norm_y * scale_factor) - (scale_factor / 2)
 
     def run_embedding_pipeline(self):
         """
-        Orchestrates the feature extraction and dimensionality reduction pipeline.
-        This version correctly re-runs reduction on cached features when parameters change.
+        Orchestrates feature extraction and dimensionality reduction.
+        If the EmbeddingViewer is in isolate mode, it will use only the visible
+        (isolated) points as input for the pipeline.
         """
-        if not self.current_data_items:
+        items_to_embed = []
+        if self.embedding_viewer.isolated_mode:
+            items_to_embed = [point.data_item for point in self.embedding_viewer.isolated_points]
+        else:
+            items_to_embed = self.current_data_items
+
+        if not items_to_embed:
             print("No items to process for embedding.")
             return
 
-        # 1. Get current parameters from the UI
+        self.annotation_viewer.clear_selection()
+        if self.annotation_viewer.isolated_mode:
+            self.annotation_viewer.show_all_annotations()
+
+        self.embedding_viewer.render_selection_from_ids(set())
+        self.update_button_states()
+
+        self.current_embedding_model_info = self.model_settings_widget.get_selected_model()
+
         embedding_params = self.embedding_settings_widget.get_embedding_parameters()
-        model_info = self.model_settings_widget.get_selected_model()  # Now returns tuple (model_name, feature_mode)
-        
-        # Unpack model info - both model_name and feature_mode are used for caching
-        if isinstance(model_info, tuple):
-            selected_model, selected_feature_mode = model_info
-            # Create a unique cache key that includes both model and feature mode
-            cache_key = f"{selected_model}_{selected_feature_mode}"
+        selected_model, selected_feature_mode = self.current_embedding_model_info
+
+        # If the model name is a path, use only its base name.
+        if os.path.sep in selected_model or '/' in selected_model:
+            sanitized_model_name = os.path.basename(selected_model)
         else:
-            selected_model = model_info
-            selected_feature_mode = "default"
-            cache_key = f"{selected_model}_{selected_feature_mode}"
-            
-        technique = embedding_params['technique']
+            sanitized_model_name = selected_model
+
+        # Replace characters that might be problematic in filenames
+        sanitized_model_name = sanitized_model_name.replace(' ', '_')
+        # Also replace the forward slash to handle "N/A"
+        sanitized_feature_mode = selected_feature_mode.replace(' ', '_').replace('/', '_')
+
+        model_key = f"{sanitized_model_name}_{sanitized_feature_mode}"
 
         QApplication.setOverrideCursor(Qt.WaitCursor)
-        progress_bar = ProgressBar(self, "Generating Embedding Visualization")
+        progress_bar = ProgressBar(self, "Processing Annotations")
         progress_bar.show()
-        
-        try:
-            # 2. Decide whether to use cached features or extract new ones
-            # Now checks both model name AND feature mode
-            if self.current_features is None or cache_key != self.current_feature_generating_model:
-                # SLOW PATH: Extract and cache new features
-                features, valid_data_items = self._extract_features(self.current_data_items, progress_bar=progress_bar)
-                
-                self.current_features = features
-                self.current_feature_generating_model = cache_key  # Store the complete cache key
-                self.current_data_items = valid_data_items
-                self.annotation_viewer.update_annotations(self.current_data_items)
-            else:
-                # FAST PATH: Use existing features
-                print("Using cached features. Skipping feature extraction.")
-                features = self.current_features
 
-            if features is None or len(features) == 0:
-                print("No valid features available. Aborting embedding.")
+        try:
+            progress_bar.set_busy_mode("Checking feature cache...")
+            cached_features, items_to_process = self.feature_store.get_features(items_to_embed, model_key)
+            print(f"Found {len(cached_features)} features in cache. Need to compute {len(items_to_process)}.")
+
+            if items_to_process:
+                newly_extracted_features, valid_items_processed = self._extract_features(items_to_process,
+                                                                                         progress_bar=progress_bar)
+                if len(newly_extracted_features) > 0:
+                    progress_bar.set_busy_mode("Saving new features to cache...")
+                    self.feature_store.add_features(valid_items_processed, newly_extracted_features, model_key)
+                    new_features_dict = {item.annotation.id: vec for item, vec in zip(valid_items_processed,
+                                                                                      newly_extracted_features)}
+                    cached_features.update(new_features_dict)
+
+            if not cached_features:
+                print("No features found or computed. Aborting.")
                 return
 
-            # 3. Run dimensionality reduction with the latest parameters
-            progress_bar.set_busy_mode(f"Running {technique} dimensionality reduction...")
+            final_feature_list = []
+            final_data_items = []
+            for item in items_to_embed:
+                if item.annotation.id in cached_features:
+                    final_feature_list.append(cached_features[item.annotation.id])
+                    final_data_items.append(item)
+
+            features = np.array(final_feature_list)
+            self.current_data_items = final_data_items
+            self.annotation_viewer.update_annotations(self.current_data_items)
+
+            progress_bar.set_busy_mode("Running dimensionality reduction...")
             embedded_features = self._run_dimensionality_reduction(features, embedding_params)
-            progress_bar.update_progress()
-            
+
             if embedded_features is None:
                 return
 
-            # 4. Update the visualization with the new 2D layout
             progress_bar.set_busy_mode("Updating visualization...")
             self._update_data_items_with_embedding(self.current_data_items, embedded_features)
-            
             self.embedding_viewer.update_embeddings(self.current_data_items)
             self.embedding_viewer.show_embedding()
             self.embedding_viewer.fit_view_to_points()
-            progress_bar.update_progress()
-            
-            print(f"Successfully generated embedding for {len(self.current_data_items)} annotations using {technique}")
-            
-        except Exception as e:
-            print(f"Error during embedding pipeline: {e}")
-            self.embedding_viewer.clear_points()
-            self.embedding_viewer.show_placeholder()
-            
+
+            # Check if confidence scores are available to enable sorting
+            _, feature_mode = self.current_embedding_model_info
+            is_predict_mode = feature_mode == "Predictions"
+            self.annotation_viewer.set_confidence_sort_availability(is_predict_mode)
+
+            # If using Predictions mode, update data items with probabilities for confidence sorting
+            if is_predict_mode:
+                for item in self.current_data_items:
+                    if item.annotation.id in cached_features:
+                        item.prediction_probabilities = cached_features[item.annotation.id]
+
+            # When a new embedding is run, any previous similarity sort becomes irrelevant
+            self.annotation_viewer.active_ordered_ids = []
+
         finally:
             QApplication.restoreOverrideCursor()
             progress_bar.finish_progress()
@@ -1956,20 +2479,15 @@ class ExplorerWindow(QMainWindow):
         """Refresh display: filter data and update annotation viewer."""
         QApplication.setOverrideCursor(Qt.WaitCursor)
         try:
-            # Get filtered data and store for potential embedding
             self.current_data_items = self.get_filtered_data_items()
-            
-            # --- MODIFIED: Invalidate the feature cache ---
-            # Since the filtered items have changed, the old features are no longer valid.
             self.current_features = None
-            
-            # Update annotation viewer with filtered data
             self.annotation_viewer.update_annotations(self.current_data_items)
-            
-            # Clear embedding viewer and show placeholder, as it is now out of sync
             self.embedding_viewer.clear_points()
             self.embedding_viewer.show_placeholder()
-            
+
+            # Reset sort options when filters change
+            self.annotation_viewer.active_ordered_ids = []
+            self.annotation_viewer.set_confidence_sort_availability(False)
         finally:
             QApplication.restoreOverrideCursor()
 
@@ -1979,89 +2497,144 @@ class ExplorerWindow(QMainWindow):
             self.annotation_viewer.apply_preview_label_to_selected(label)
             self.update_button_states()
 
+    def delete_data_items(self, data_items_to_delete):
+        """
+        Permanently deletes a list of data items and their associated annotations
+        and visual components from the explorer and the main application.
+        """
+        if not data_items_to_delete:
+            return
+
+        print(f"Permanently deleting {len(data_items_to_delete)} item(s).")
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        try:
+            deleted_ann_ids = {item.annotation.id for item in data_items_to_delete}
+            annotations_to_delete_from_main_app = [item.annotation for item in data_items_to_delete]
+
+            # 1. Delete from the main application's data store
+            self.annotation_window.delete_annotations(annotations_to_delete_from_main_app)
+
+            # 2. Remove from Explorer's internal data structures
+            self.current_data_items = [
+                item for item in self.current_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]
+
+            # 3. Remove from AnnotationViewer
+            blocker = QSignalBlocker(self.annotation_viewer)  # Block signals during mass removal
+            for ann_id in deleted_ann_ids:
+                if ann_id in self.annotation_viewer.annotation_widgets_by_id:
+                    widget = self.annotation_viewer.annotation_widgets_by_id.pop(ann_id)
+                    if widget in self.annotation_viewer.selected_widgets:
+                        self.annotation_viewer.selected_widgets.remove(widget)
+                    widget.setParent(None)
+                    widget.deleteLater()
+            blocker.unblock()
+            self.annotation_viewer.recalculate_widget_positions()
+
+            # 4. Remove from EmbeddingViewer
+            blocker = QSignalBlocker(self.embedding_viewer.graphics_scene)
+            for ann_id in deleted_ann_ids:
+                if ann_id in self.embedding_viewer.points_by_id:
+                    point = self.embedding_viewer.points_by_id.pop(ann_id)
+                    self.embedding_viewer.graphics_scene.removeItem(point)
+            blocker.unblock()
+            self.embedding_viewer.on_selection_changed()  # Trigger update of selection state
+
+            # 5. Update UI
+            self.update_label_window_selection()
+            self.update_button_states()
+
+            # 6. Refresh main window annotations list
+            affected_images = {ann.image_path for ann in annotations_to_delete_from_main_app}
+            for image_path in affected_images:
+                self.image_window.update_image_annotations(image_path)
+            self.annotation_window.load_annotations()
+
+        except Exception as e:
+            print(f"Error during item deletion: {e}")
+        finally:
+            QApplication.restoreOverrideCursor()
+
     def clear_preview_changes(self):
-        """Clear all preview changes and revert to original labels."""
+        """
+        Clears all preview changes in the annotation viewer and updates tooltips.
+        """
         if hasattr(self, 'annotation_viewer'):
             self.annotation_viewer.clear_preview_states()
-            self.update_button_states()
-            print("Cleared all preview changes")
+
+            # After reverting, tooltips need to be updated to reflect original labels
+            for widget in self.annotation_viewer.annotation_widgets_by_id.values():
+                widget.update_tooltip()
+            for point in self.embedding_viewer.points_by_id.values():
+                point.update_tooltip()
+
+        # After reverting all changes, update the button states
+        self.update_button_states()
+        print("Cleared all pending changes.")
 
     def update_button_states(self):
-        """Update the state of Clear Preview and Apply buttons."""
-        has_changes = (hasattr(self, 'annotation_viewer') and self.annotation_viewer.has_preview_changes())
-        
+        """Update the state of Clear Preview, Apply, and Find Similar buttons."""
+        has_changes = self.annotation_viewer.has_preview_changes()
         self.clear_preview_button.setEnabled(has_changes)
         self.apply_button.setEnabled(has_changes)
         
+        # Update tooltips with a summary of changes
         summary = self.annotation_viewer.get_preview_changes_summary()
         self.clear_preview_button.setToolTip(f"Clear all preview changes - {summary}")
         self.apply_button.setToolTip(f"Apply changes - {summary}")
 
+        # Logic for the "Find Similar" button
+        selection_exists = bool(self.annotation_viewer.selected_widgets)
+        embedding_exists = bool(self.embedding_viewer.points_by_id) and self.current_embedding_model_info is not None
+        self.annotation_viewer.find_similar_button.setEnabled(selection_exists and embedding_exists)
+
     def apply(self):
-        """Apply any modifications to the actual annotations."""
-        # Make cursor busy
+        """
+        Apply all pending label modifications to the main application's data.
+        """
         QApplication.setOverrideCursor(Qt.WaitCursor)
-        
         try:
-            applied_annotations = self.annotation_viewer.apply_preview_changes_permanently()
-            
-            if applied_annotations:
-                # Find which data items were affected and tell their visual components to update
-                changed_ids = {ann.id for ann in applied_annotations}
-                for item in self.current_data_items:
-                    if item.annotation.id in changed_ids:
-                        # Update annotation widget in the grid
-                        widget = self.annotation_viewer.annotation_widgets_by_id.get(item.annotation.id)
-                        if widget:
-                            widget.update() # Repaint to show new permanent color
-
-                        # Update point in the embedding viewer
-                        point = self.embedding_viewer.points_by_id.get(item.annotation.id)
-                        if point:
-                            point.update() # Repaint to show new permanent color
-
-                # Update sorting if currently sorting by label
-                if self.annotation_viewer.sort_combo.currentText() == "Label":
-                    self.annotation_viewer.recalculate_widget_positions()
-
-                # Update the main application's data
-                affected_images = {ann.image_path for ann in applied_annotations}
-                for image_path in affected_images:
-                    self.image_window.update_image_annotations(image_path)
-                self.annotation_window.load_annotations()
-                
-                # Clear selection and button states
-                self.annotation_viewer.clear_selection()
-                self.embedding_viewer.render_selection_from_ids(set()) # Clear embedding selection
-                self.update_button_states()
-                
-                print(f"Applied changes to {len(applied_annotations)} annotation(s)")
-            else:
-                print("No preview changes to apply")
+            # --- 1. Process Label Changes ---
+            applied_label_changes = []
+            # Iterate over all current data items
+            for item in self.current_data_items:
+                if item.apply_preview_permanently():
+                    applied_label_changes.append(item.annotation)
+
+            # --- 2. Update UI if any changes were made ---
+            if not applied_label_changes:
+                print("No pending changes to apply.")
+                return
+
+            # Update the main application's data and UI
+            affected_images = {ann.image_path for ann in applied_label_changes}
+            for image_path in affected_images:
+                self.image_window.update_image_annotations(image_path)
+            self.annotation_window.load_annotations()
+
+            # Refresh the annotation viewer since its underlying data has changed
+            self.annotation_viewer.update_annotations(self.current_data_items)
+
+            # Reset selections and button states
+            self.embedding_viewer.render_selection_from_ids(set())
+            self.update_label_window_selection()
+            self.update_button_states()
+
+            print("Applied changes successfully.")
 
         except Exception as e:
             print(f"Error applying modifications: {e}")
         finally:
-            # Restore cursor
             QApplication.restoreOverrideCursor()
             
     def _cleanup_resources(self):
-        """
-        Clean up heavy resources like the loaded model and clear GPU cache.
-        This is called when the window is closed to free up memory.
-        """
-        print("Cleaning up Explorer resources...")
-        
-        # Reset model and feature caches
-        self.imgsz = 128
+        """Clean up resources."""
         self.loaded_model = None
         self.model_path = ""
         self.current_features = None
         self.current_feature_generating_model = ""
-        
-        # Clear CUDA cache if available to free up GPU memory
         if torch.cuda.is_available():
-            print("Clearing CUDA cache.")
-            torch.cuda.empty_cache()
-            
-        print("Cleanup complete.")
+            torch.cuda.empty_cache()