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

coralnet_toolbox/SeeAnything/QtDeployGenerator.py

@@ -2,24 +2,24 @@ import warnings
 
 import os
 import gc
-import json
-import copy
 
 import numpy as np
+from sklearn.decomposition import PCA
 
 import torch
 from torch.cuda import empty_cache
 
+import pyqtgraph as pg
+from pyqtgraph.Qt import QtGui
+
 from ultralytics import YOLOE
 from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
-from ultralytics.models.yolo.yoloe import YOLOEVPDetectPredictor
 
 from PyQt5.QtCore import Qt
-from PyQt5.QtGui import QColor
-from PyQt5.QtWidgets import (QMessageBox, QCheckBox, QVBoxLayout, QApplication,
+from PyQt5.QtWidgets import (QMessageBox, QVBoxLayout, QApplication, QFileDialog,
                              QLabel, QDialog, QDialogButtonBox, QGroupBox, QLineEdit,
                              QFormLayout, QComboBox, QSpinBox, QSlider, QPushButton,
-                             QHBoxLayout, QWidget, QFileDialog)
+                             QHBoxLayout)
 
 from coralnet_toolbox.Annotations.QtPolygonAnnotation import PolygonAnnotation
 from coralnet_toolbox.Annotations.QtRectangleAnnotation import RectangleAnnotation
@@ -79,10 +79,18 @@ class DeployGeneratorDialog(QDialog):
         self.class_mapping = {}
 
         # Reference image and label
-        self.source_images = []
-        self.source_label = None
-        # Target images
-        self.target_images = []
+        self.reference_label = None
+        self.reference_image_paths = []
+        
+        # Visual Prompting Encoding (VPE) - legacy single tensor variable
+        self.vpe_path = None
+        self.vpe = None
+        
+        # New separate VPE collections
+        self.imported_vpes = []  # VPEs loaded from file
+        self.reference_vpes = []  # VPEs created from reference images
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         # Main vertical layout for the dialog
         self.layout = QVBoxLayout(self)
@@ -110,7 +118,7 @@ class DeployGeneratorDialog(QDialog):
         self.setup_status_layout()
         
         # Add layouts to the right panel
-        self.setup_source_layout()
+        self.setup_reference_layout()
 
         # # Add a full ImageWindow instance for target image selection
         self.image_selection_window = ImageWindow(self.main_window)
@@ -158,6 +166,9 @@ class DeployGeneratorDialog(QDialog):
         iw.search_bar_images.setEnabled(False)
         iw.search_bar_labels.setEnabled(False)
         iw.top_k_combo.setEnabled(False)
+        
+        # Hide the "Current" label as it is not applicable in this dialog
+        iw.current_image_index_label.hide()
 
         # Set Top-K to Top1
         iw.top_k_combo.setCurrentText("Top1")
@@ -190,7 +201,7 @@ class DeployGeneratorDialog(QDialog):
         self.sync_image_window()
         # This now populates the dropdown, restores the last selection,
         # and then manually triggers the image filtering.
-        self.update_source_labels()
+        self.update_reference_labels()
 
     def sync_image_window(self):
         """
@@ -219,14 +230,23 @@ class DeployGeneratorDialog(QDialog):
         annotation that has BOTH the selected label AND a valid type (Polygon or Rectangle).
         This uses the fast, pre-computed cache for performance.
         """
-        source_label = self.source_label_combo_box.currentData()
-        source_label_text = self.source_label_combo_box.currentText()
+        # Persist the user's current highlights from the table model before filtering.
+        # This ensures that if the user highlights items and then changes the filter,
+        # their selection is not lost.
+        current_highlights = self.image_selection_window.table_model.get_highlighted_paths()
+        if current_highlights:
+            self.reference_image_paths = current_highlights
+
+        reference_label = self.reference_label_combo_box.currentData()
+        reference_label_text = self.reference_label_combo_box.currentText()
 
         # Store the last selected label for a better user experience on re-opening.
-        if source_label_text:
-            self.last_selected_label_code = source_label_text
+        if reference_label_text:
+            self.last_selected_label_code = reference_label_text
+            # Also store the reference label object itself
+            self.reference_label = reference_label
 
-        if not source_label:
+        if not reference_label:
             # If no label is selected (e.g., during initialization), show an empty list.
             self.image_selection_window.table_model.set_filtered_paths([])
             return
@@ -235,7 +255,7 @@ class DeployGeneratorDialog(QDialog):
         final_filtered_paths = []
         
         valid_types = {"RectangleAnnotation", "PolygonAnnotation"}
-        selected_label_code = source_label.short_label_code
+        selected_label_code = reference_label.short_label_code
 
         # Loop through paths and check the pre-computed map on each raster
         for path in all_paths:
@@ -257,40 +277,54 @@ class DeployGeneratorDialog(QDialog):
         # Directly set the filtered list in the table model.
         self.image_selection_window.table_model.set_filtered_paths(final_filtered_paths)
         
+        # Try to preserve any previous selections
+        if hasattr(self, 'reference_image_paths') and self.reference_image_paths:
+            # Find which of our previously selected paths are still in the filtered list
+            valid_selections = [p for p in self.reference_image_paths if p in final_filtered_paths]
+            if valid_selections:
+                # Highlight previously selected paths that are still valid
+                self.image_selection_window.table_model.set_highlighted_paths(valid_selections)
+
+        # After filtering, update all labels with the correct counts.
+        dialog_iw = self.image_selection_window
+        dialog_iw.update_image_count_label(len(final_filtered_paths)) # Set "Total" to filtered count
+        dialog_iw.update_current_image_index_label()
+        dialog_iw.update_highlighted_count_label()
+                
     def accept(self):
         """
         Validate selections and store them before closing the dialog.
+        A prediction is valid if a model and label are selected, and the user
+        has provided either reference images or an imported VPE file.
         """
         if not self.loaded_model:
             QMessageBox.warning(self, 
                                 "No Model", 
                                 "A model must be loaded before running predictions.")
-            super().reject()
             return
 
-        current_label = self.source_label_combo_box.currentData()
-        if not current_label:
+        # Set reference label from combo box
+        self.reference_label = self.reference_label_combo_box.currentData()
+        if not self.reference_label:
             QMessageBox.warning(self, 
-                                "No Source Label", 
-                                "A source label must be selected.")
-            super().reject()
+                                "No Reference Label", 
+                                "A reference label must be selected.")
             return
 
-        # Get highlighted paths from our internal image window to use as targets
-        highlighted_images = self.image_selection_window.table_model.get_highlighted_paths()
+        # Stash the current UI selection before validating.
+        self.update_stashed_references_from_ui()
+
+        # Check for a valid VPE source using the now-stashed list.
+        has_reference_images = bool(self.reference_image_paths)
+        has_imported_vpes = bool(self.imported_vpes)
 
-        if not highlighted_images:
+        if not has_reference_images and not has_imported_vpes:
             QMessageBox.warning(self, 
-                                "No Target Images", 
-                                "You must highlight at least one image in the list to process.")
-            super().reject()
+                                "No VPE Source Provided", 
+                                "You must highlight at least one reference image or load a VPE file to proceed.")
             return
 
-        # Store the selections for the caller to use after the dialog closes.
-        self.source_label = current_label
-        self.target_images = highlighted_images
-
-        # Do not call self.predict here; just close the dialog and let the caller handle prediction
+        # If validation passes, close the dialog.
         super().accept()
 
     def setup_info_layout(self):
@@ -317,19 +351,19 @@ class DeployGeneratorDialog(QDialog):
         Setup the models layout with a simple model selection combo box (no tabs).
         """
         group_box = QGroupBox("Model Selection")
-        layout = QVBoxLayout()
+        layout = QFormLayout()
 
         self.model_combo = QComboBox()
         self.model_combo.setEditable(True)
 
         # Define available models (keep the existing dictionary)
         self.models = [
-            "yoloe-v8s-seg.pt",
-            "yoloe-v8m-seg.pt",
-            "yoloe-v8l-seg.pt",
-            "yoloe-11s-seg.pt",
-            "yoloe-11m-seg.pt",
-            "yoloe-11l-seg.pt",
+            'yoloe-v8s-seg.pt',
+            'yoloe-v8m-seg.pt',
+            'yoloe-v8l-seg.pt',
+            'yoloe-11s-seg.pt',
+            'yoloe-11m-seg.pt',
+            'yoloe-11l-seg.pt',
         ]
 
         # Add all models to combo box
@@ -337,10 +371,19 @@ class DeployGeneratorDialog(QDialog):
             self.model_combo.addItem(model_name)
         
         # Set the default model
-        self.model_combo.setCurrentText("yoloe-v8s-seg.pt")
+        self.model_combo.setCurrentIndex(self.models.index('yoloe-v8s-seg.pt'))
+        # Create a layout for the model selection
+        layout.addRow(QLabel("Models:"), self.model_combo)
+
+        # Add custom vpe file selection
+        self.vpe_path_edit = QLineEdit()
+        browse_button = QPushButton("Browse...")
+        browse_button.clicked.connect(self.browse_vpe_file)
 
-        layout.addWidget(QLabel("Select Model:"))
-        layout.addWidget(self.model_combo)
+        vpe_path_layout = QHBoxLayout()
+        vpe_path_layout.addWidget(self.vpe_path_edit)
+        vpe_path_layout.addWidget(browse_button)
+        layout.addRow("Custom VPE:", vpe_path_layout)
 
         group_box.setLayout(layout)
         self.left_panel.addWidget(group_box)  # Add to left panel
@@ -373,7 +416,7 @@ class DeployGeneratorDialog(QDialog):
 
         # Image size control
         self.imgsz_spinbox = QSpinBox()
-        self.imgsz_spinbox.setRange(512, 65536)
+        self.imgsz_spinbox.setRange(1024, 65536)
         self.imgsz_spinbox.setSingleStep(1024)
         self.imgsz_spinbox.setValue(self.imgsz)
         layout.addRow("Image Size (imgsz):", self.imgsz_spinbox)
@@ -445,17 +488,38 @@ class DeployGeneratorDialog(QDialog):
         Setup action buttons in a group box.
         """
         group_box = QGroupBox("Actions")
-        layout = QHBoxLayout()
+        main_layout = QVBoxLayout()
 
+        # First row: Load and Deactivate buttons side by side
+        button_row = QHBoxLayout()
         load_button = QPushButton("Load Model")
         load_button.clicked.connect(self.load_model)
-        layout.addWidget(load_button)
+        button_row.addWidget(load_button)
 
         deactivate_button = QPushButton("Deactivate Model")
         deactivate_button.clicked.connect(self.deactivate_model)
-        layout.addWidget(deactivate_button)
+        button_row.addWidget(deactivate_button)
 
-        group_box.setLayout(layout)
+        main_layout.addLayout(button_row)
+
+        # Second row: VPE action buttons
+        vpe_row = QHBoxLayout()
+        
+        generate_vpe_button = QPushButton("Generate VPEs")
+        generate_vpe_button.clicked.connect(self.generate_vpes_from_references)
+        vpe_row.addWidget(generate_vpe_button)
+
+        save_vpe_button = QPushButton("Save VPE")
+        save_vpe_button.clicked.connect(self.save_vpe)
+        vpe_row.addWidget(save_vpe_button)
+
+        show_vpe_button = QPushButton("Show VPE")
+        show_vpe_button.clicked.connect(self.show_vpe)
+        vpe_row.addWidget(show_vpe_button)
+
+        main_layout.addLayout(vpe_row)
+
+        group_box.setLayout(main_layout)
         self.left_panel.addWidget(group_box)  # Add to left panel
 
     def setup_status_layout(self):
@@ -471,18 +535,23 @@ class DeployGeneratorDialog(QDialog):
         group_box.setLayout(layout)
         self.left_panel.addWidget(group_box)  # Add to left panel
 
-    def setup_source_layout(self):
+    def setup_reference_layout(self):
         """
-        Set up the layout with source label selection.
-        The source image is implicitly the currently active image.
+        Set up the layout with reference label selection.
+        The reference image is implicitly the currently active image.
         """
-        group_box = QGroupBox("Reference Label")
+        group_box = QGroupBox("Reference")
         layout = QFormLayout()
 
-        # Create the source label combo box
-        self.source_label_combo_box = QComboBox()
-        self.source_label_combo_box.currentIndexChanged.connect(self.filter_images_by_label_and_type)
-        layout.addRow("Reference Label:", self.source_label_combo_box)
+        # Create the reference label combo box
+        self.reference_label_combo_box = QComboBox()
+        self.reference_label_combo_box.currentIndexChanged.connect(self.filter_images_by_label_and_type)
+        layout.addRow("Reference Label:", self.reference_label_combo_box)
+        
+        # Create a Reference model combobox (VPE, Images)
+        self.reference_method_combo_box = QComboBox()
+        self.reference_method_combo_box.addItems(["VPE", "Images"])
+        layout.addRow("Reference Method:", self.reference_method_combo_box)
 
         group_box.setLayout(layout)
         self.right_panel.addWidget(group_box)  # Add to right panel
@@ -543,6 +612,10 @@ class DeployGeneratorDialog(QDialog):
         self.area_thresh_max = max_val / 100.0
         self.main_window.update_area_thresh(self.area_thresh_min, self.area_thresh_max)
         self.area_threshold_label.setText(f"{self.area_thresh_min:.2f} - {self.area_thresh_max:.2f}")
+        
+    def update_stashed_references_from_ui(self):
+        """Updates the internal reference path list from the current UI selection."""
+        self.reference_image_paths = self.image_selection_window.table_model.get_highlighted_paths()
 
     def get_max_detections(self):
         """Get the maximum number of detections to return."""
@@ -603,54 +676,44 @@ class DeployGeneratorDialog(QDialog):
             if self.loaded_model:
                 self.deactivate_model()
 
-    def update_source_labels(self):
+    def update_reference_labels(self):
         """
-        Updates the source label combo box with labels that are associated with
-        valid reference annotations (Polygons or Rectangles), using the fast cache.
+        Updates the reference label combo box with ALL available project labels.
+        This dropdown now serves as the "Output Label" for all predictions.
+        The "Review" label with id "-1" is excluded.
         """
-        self.source_label_combo_box.blockSignals(True)
+        self.reference_label_combo_box.blockSignals(True)
         
         try:
-            self.source_label_combo_box.clear()
-
-            dialog_manager = self.image_selection_window.raster_manager
-            valid_types = {"RectangleAnnotation", "PolygonAnnotation"}
-            valid_labels = set()  # This will store the full Label objects
-
-            # Create a lookup map to get full label objects from their codes
-            all_project_labels = {lbl.short_label_code: lbl for lbl in self.main_window.label_window.labels}
-
-            # Use the cached data to find all labels that have valid reference types.
-            for raster in dialog_manager.rasters.values():
-                # raster.label_to_types_map is like: {'coral': {'Point'}, 'rock': {'Polygon'}}
-                for label_code, types_for_label in raster.label_to_types_map.items():
-                    # Check if the set of types for this specific label
-                    # has any overlap with our valid reference types.
-                    if not valid_types.isdisjoint(types_for_label):
-                        # This label is a valid reference label.
-                        # Add its full Label object to our set of valid labels.
-                        if label_code in all_project_labels:
-                            valid_labels.add(all_project_labels[label_code])
+            self.reference_label_combo_box.clear()
+
+            # Get all labels from the main label window
+            all_project_labels = self.main_window.label_window.labels
+
+            # Filter out the special "Review" label and create a list of valid labels
+            valid_labels = [
+                label_obj for label_obj in all_project_labels
+                if not (label_obj.short_label_code == "Review" and str(label_obj.id) == "-1")
+            ]
 
             # Add the valid labels to the combo box, sorted alphabetically.
-            sorted_valid_labels = sorted(list(valid_labels), key=lambda x: x.short_label_code)
+            sorted_valid_labels = sorted(valid_labels, key=lambda x: x.short_label_code)
             for label_obj in sorted_valid_labels:
-                self.source_label_combo_box.addItem(label_obj.short_label_code, label_obj)
+                self.reference_label_combo_box.addItem(label_obj.short_label_code, label_obj)
 
             # Restore the last selected label if it's still present in the list.
             if self.last_selected_label_code:
-                index = self.source_label_combo_box.findText(self.last_selected_label_code)
+                index = self.reference_label_combo_box.findText(self.last_selected_label_code)
                 if index != -1:
-                    self.source_label_combo_box.setCurrentIndex(index)
+                    self.reference_label_combo_box.setCurrentIndex(index)
         finally:
-            self.source_label_combo_box.blockSignals(False)
+            self.reference_label_combo_box.blockSignals(False)
         
-        # Manually trigger the filtering now that the combo box is stable.
+        # Manually trigger the image filtering now that the combo box is stable.
+        # This will still filter the image list to help find references if needed.
         self.filter_images_by_label_and_type()
 
-        return True
-
-    def get_source_annotations(self, reference_label, reference_image_path):
+    def get_reference_annotations(self, reference_label, reference_image_path):
         """
         Return a list of bboxes and masks for a specific image
         belonging to the selected label.
@@ -666,22 +729,161 @@ class DeployGeneratorDialog(QDialog):
         annotations = self.annotation_window.get_image_annotations(reference_image_path)
 
         # Filter annotations by the provided label
-        source_bboxes = []
-        source_masks = []
+        reference_bboxes = []
+        reference_masks = []
         for annotation in annotations:
             if annotation.label.short_label_code == reference_label.short_label_code:
                 if isinstance(annotation, (PolygonAnnotation, RectangleAnnotation)):
                     bbox = annotation.cropped_bbox
-                    source_bboxes.append(bbox)
+                    reference_bboxes.append(bbox)
                     if isinstance(annotation, PolygonAnnotation):
                         points = np.array([[p.x(), p.y()] for p in annotation.points])
-                        source_masks.append(points)
+                        reference_masks.append(points)
                     elif isinstance(annotation, RectangleAnnotation):
                         x1, y1, x2, y2 = bbox
                         rect_points = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])
-                        source_masks.append(rect_points)
+                        reference_masks.append(rect_points)
 
-        return np.array(source_bboxes), source_masks
+        return np.array(reference_bboxes), reference_masks
+    
+    def browse_vpe_file(self):
+        """
+        Open a file dialog to browse for a VPE file and load it.
+        Stores imported VPEs separately from reference-generated VPEs.
+        """
+        file_path, _ = QFileDialog.getOpenFileName(
+            self,
+            "Select Visual Prompt Encoding (VPE) File",
+            "",
+            "VPE Files (*.pt);;All Files (*)"
+        )
+        
+        if not file_path:
+            return
+            
+        self.vpe_path_edit.setText(file_path)
+        self.vpe_path = file_path
+        
+        try:
+            # Load the VPE file
+            loaded_data = torch.load(file_path)
+            
+            # TODO Move tensors to the appropriate device
+            # device = self.main_window.device
+            
+            # Check format type and handle appropriately
+            if isinstance(loaded_data, list):
+                # New format: list of VPE tensors
+                self.imported_vpes = [vpe.to(self.device) for vpe in loaded_data]
+                vpe_count = len(self.imported_vpes)
+                self.status_bar.setText(f"Loaded {vpe_count} VPE tensors from file")
+                
+            elif isinstance(loaded_data, torch.Tensor):
+                # Legacy format: single tensor - convert to list for consistency
+                loaded_vpe = loaded_data.to(self.device)
+                # Store as a single-item list
+                self.imported_vpes = [loaded_vpe]
+                self.status_bar.setText("Loaded 1 VPE tensor from file (legacy format)")
+                
+            else:
+                # Invalid format
+                self.imported_vpes = []
+                self.status_bar.setText("Invalid VPE file format")
+                QMessageBox.warning(
+                    self, 
+                    "Invalid VPE", 
+                    "The file does not appear to be a valid VPE format."
+                )
+                # Clear the VPE path edit field
+                self.vpe_path_edit.clear()
+                    
+            # For backward compatibility - set self.vpe to the average of imported VPEs
+            # This ensures older code paths still work
+            if self.imported_vpes:
+                combined_vpe = torch.cat(self.imported_vpes).mean(dim=0, keepdim=True)
+                self.vpe = torch.nn.functional.normalize(combined_vpe, p=2, dim=-1)
+                
+        except Exception as e:
+            self.imported_vpes = []
+            self.vpe = None
+            self.status_bar.setText(f"Error loading VPE: {str(e)}")
+            QMessageBox.critical(
+                self, 
+                "Error Loading VPE", 
+                f"Failed to load VPE file: {str(e)}"
+            )
+            
+    def save_vpe(self):
+        """
+        Saves the combined collection of VPEs (imported and pre-generated from references) to disk.
+        """
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        
+        try:
+            # Create a list to hold all VPEs to be saved
+            all_vpes = []
+            
+            # Add imported VPEs if available
+            if self.imported_vpes:
+                all_vpes.extend(self.imported_vpes)
+            
+            # Add pre-generated reference VPEs if available
+            if self.reference_vpes:
+                all_vpes.extend(self.reference_vpes)
+            
+            # Check if we have any VPEs to save
+            if not all_vpes:
+                QApplication.restoreOverrideCursor()
+                QMessageBox.warning(
+                    self,
+                    "No VPEs Available",
+                    "No VPEs available to save. "
+                    "Please either load a VPE file or generate VPEs from reference images first."
+                )
+                return
+            
+            QApplication.restoreOverrideCursor()
+            
+            file_path, _ = QFileDialog.getSaveFileName(
+                self,
+                "Save VPE Collection",
+                "",
+                "PyTorch Tensor (*.pt);;All Files (*)"
+            )
+            
+            if not file_path:
+                return
+            
+            QApplication.setOverrideCursor(Qt.WaitCursor)
+            
+            if not file_path.endswith('.pt'):
+                file_path += '.pt'
+            
+            vpe_list_cpu = [vpe.cpu() for vpe in all_vpes]
+            
+            torch.save(vpe_list_cpu, file_path)
+            
+            self.status_bar.setText(f"Saved {len(all_vpes)} VPE tensors to {os.path.basename(file_path)}")
+            
+            QApplication.restoreOverrideCursor()
+            QMessageBox.information(
+                self,
+                "VPE Saved",
+                f"Saved {len(all_vpes)} VPE tensors to {file_path}"
+            )
+            
+        except Exception as e:
+            QApplication.restoreOverrideCursor()
+            QMessageBox.critical(
+                self,
+                "Error Saving VPE",
+                f"Failed to save VPE: {str(e)}"
+            )
+        finally:
+            try:
+                QApplication.restoreOverrideCursor()
+            except:
+                pass
 
     def load_model(self):
         """
@@ -692,40 +894,32 @@ class DeployGeneratorDialog(QDialog):
         progress_bar.show()
 
         try:
-            # Get selected model path and download weights if needed
-            self.model_path = self.model_combo.currentText()
-
-            # Load model using registry
-            self.loaded_model = YOLOE(self.model_path).to(self.main_window.device)
-
-            # Create a dummy visual dictionary
-            visuals = dict(
-                bboxes=np.array(
-                    [
-                        [120, 425, 160, 445],
-                    ],
-                ),
-                cls=np.array(
-                    np.zeros(1),
-                ),
-            )
+            # Load the model using reload_model method
+            self.reload_model()
 
-            # Run a dummy prediction to load the model
-            self.loaded_model.predict(
-                np.zeros((640, 640, 3), dtype=np.uint8),
-                visual_prompts=visuals.copy(),
-                predictor=YOLOEVPDetectPredictor,
-                imgsz=640,
-                conf=0.99,
-            )
+            # Calculate total number of VPEs from both sources
+            total_vpes = len(self.imported_vpes) + len(self.reference_vpes)
+            
+            if total_vpes > 0:
+                if self.imported_vpes and self.reference_vpes:
+                    message = f"Model loaded with {len(self.imported_vpes)} imported VPEs "
+                    message += f"and {len(self.reference_vpes)} reference VPEs"
+                elif self.imported_vpes:
+                    message = f"Model loaded with {len(self.imported_vpes)} imported VPEs"
+                else:
+                    message = f"Model loaded with {len(self.reference_vpes)} reference VPEs"
+                    
+                self.status_bar.setText(message)
+            else:
+                message = "Model loaded with default VPE"
+                self.status_bar.setText("Model loaded with default VPE")
 
+            # Finish progress bar
             progress_bar.finish_progress()
-            self.status_bar.setText("Model loaded")
-            QMessageBox.information(self.annotation_window, 
-                                    "Model Loaded", 
-                                    "Model loaded successfully")
+            QMessageBox.information(self.annotation_window, "Model Loaded", message)
 
         except Exception as e:
+            self.loaded_model = None
             QMessageBox.critical(self.annotation_window, 
                                  "Error Loading Model", 
                                  f"Error loading model: {e}")
@@ -737,7 +931,52 @@ class DeployGeneratorDialog(QDialog):
             progress_bar.stop_progress()
             progress_bar.close()
             progress_bar = None
+            
+    def reload_model(self):
+        """
+        Subset of the load_model method. This is needed when additional 
+        reference images and annotations (i.e., VPEs) are added (we have 
+        to re-load the model each time).
+        
+        This method also ensures that we stash the currently highlighted reference
+        image paths before reloading, so they're available for predictions
+        even if the user switches the active image in the main window.
+        """        
+        self.loaded_model = None
+        
+        # Get selected model path and download weights if needed
+        self.model_path = self.model_combo.currentText()
+
+        # Load model using registry
+        self.loaded_model = YOLOE(self.model_path, verbose=False).to(self.device)  # TODO
+
+        # Create a dummy visual dictionary for standard model loading
+        visual_prompts = dict(
+            bboxes=np.array(
+                [
+                    [120, 425, 160, 445],  # Random box
+                ],
+            ),
+            cls=np.array(
+                np.zeros(1),
+            ),
+        )
+
+        # Run a dummy prediction to load the model
+        self.loaded_model.predict(
+            np.zeros((640, 640, 3), dtype=np.uint8),
+            visual_prompts=visual_prompts.copy(),  # This needs to happen to properly initialize the predictor
+            predictor=YOLOEVPSegPredictor,  # This also needs to be SegPredictor, no matter what
+            imgsz=640,
+            conf=0.99,
+        )
 
+        # If a VPE file was loaded, use it with the model after the dummy prediction
+        if self.vpe is not None and isinstance(self.vpe, torch.Tensor):
+            # Directly set the final tensor as the prompt for the predictor
+            self.loaded_model.is_fused = lambda: False
+            self.loaded_model.set_classes(["object0"], self.vpe)
+            
     def predict(self, image_paths=None):
         """
         Make predictions on the given image paths using the loaded model.
@@ -745,11 +984,11 @@ class DeployGeneratorDialog(QDialog):
         Args:
             image_paths: List of image paths to process. If None, uses the current image.
         """
-        if not self.loaded_model or not self.source_label:
+        if not self.loaded_model or not self.reference_label:
             return
         
         # Update class mapping with the selected reference label
-        self.class_mapping = {0: self.source_label}
+        self.class_mapping = {0: self.reference_label}
 
         # Create a results processor
         results_processor = ResultsProcessor(
@@ -809,32 +1048,33 @@ class DeployGeneratorDialog(QDialog):
 
         return work_areas_data
     
-    def _apply_model(self, inputs):
-        """
-        Apply the model to the target inputs, using each highlighted source
-        image as an individual reference for a separate prediction run.
+    def _get_references(self):
         """
-        # Update the model with user parameters
-        self.loaded_model.conf = self.main_window.get_uncertainty_thresh()
-        self.loaded_model.iou = self.main_window.get_iou_thresh()
-        self.loaded_model.max_det = self.get_max_detections()
+        Get the reference annotations using the stashed list of reference images
+        that was saved when the user accepted the dialog.
         
-        # NOTE: self.target_images contains the reference images highlighted in the dialog
-        reference_image_paths = self.target_images
-
-        if not reference_image_paths:
-            QMessageBox.warning(self, 
-                                "No Reference Images", 
-                                "You must highlight at least one reference image.")
-            return []
-
-        # Get the selected reference label from the stored variable
-        source_label = self.source_label
+        Returns:
+            dict: Dictionary mapping image paths to annotation data, or empty dict if no valid references.
+        """
+        # Use the "stashed" list of paths. Do NOT query the table_model again,
+        # as the UI's highlight state may have been cleared by other actions.
+        reference_paths = self.reference_image_paths
+
+        if not reference_paths:
+            print("No reference image paths were stashed to use for prediction.")
+            return {}
+
+        # Get the reference label that was also stashed
+        reference_label = self.reference_label
+        if not reference_label:
+            # This check is a safeguard; the accept() method should prevent this.
+            print("No reference label was selected.")
+            return {}
         
-        # Create a dictionary of reference annotations, with image path as the key.
+        # Create a dictionary of reference annotations from the stashed paths
         reference_annotations_dict = {}
-        for path in reference_image_paths:
-            bboxes, masks = self.get_source_annotations(source_label, path)
+        for path in reference_paths:
+            bboxes, masks = self.get_reference_annotations(reference_label, path)
             if bboxes.size > 0:
                 reference_annotations_dict[path] = {
                     'bboxes': bboxes,
@@ -842,37 +1082,48 @@ class DeployGeneratorDialog(QDialog):
                     'cls': np.zeros(len(bboxes))
                 }
 
-        # Set the task
-        self.task = self.use_task_dropdown.currentText()
-        predictor = YOLOEVPSegPredictor if self.task == "segment" else YOLOEVPDetectPredictor
+        return reference_annotations_dict
+
+    def _apply_model_using_images(self, inputs, reference_dict):
+        """
+        Apply the model using the provided images and reference annotations (dict). This method
+        loops through each reference image using its annotations; we then aggregate
+        all the results together. Less efficient, but potentially more accurate.
+
+        Args:
+            inputs (list): List of input images.
+            reference_dict (dict): Dictionary containing reference annotations for each image.
 
+        Returns:
+            list: List of prediction results.
+        """
         # Create a progress bar for iterating through reference images
         QApplication.setOverrideCursor(Qt.WaitCursor)
         progress_bar = ProgressBar(self.annotation_window, title="Making Predictions per Reference")
         progress_bar.show()
-        progress_bar.start_progress(len(reference_annotations_dict))
-        
+        progress_bar.start_progress(len(reference_dict))
+
         results_list = []
         # The 'inputs' list contains work areas from the single target image.
         # We will predict on the first work area/full image.
         input_image = inputs[0] 
 
         # Iterate through each reference image and its annotations
-        for ref_path, ref_annotations in reference_annotations_dict.items():
+        for ref_path, ref_annotations in reference_dict.items():
             # The 'refer_image' parameter is the path to the current reference image
             # The 'visual_prompts' are the annotations from that same reference image
-            visuals = {
+            visual_prompts = {
                 'bboxes': ref_annotations['bboxes'],
                 'cls': ref_annotations['cls'],
             }
             if self.task == 'segment':
-                visuals['masks'] = ref_annotations['masks']
+                visual_prompts['masks'] = ref_annotations['masks']
 
             # Make predictions on the target using the current reference
             results = self.loaded_model.predict(input_image,
                                                 refer_image=ref_path,
-                                                visual_prompts=visuals,
-                                                predictor=predictor,
+                                                visual_prompts=visual_prompts,
+                                                predictor=YOLOEVPSegPredictor,  # TODO This is necessary here?
                                                 imgsz=self.imgsz_spinbox.value(),
                                                 conf=self.main_window.get_uncertainty_thresh(),
                                                 iou=self.main_window.get_iou_thresh(),
@@ -904,6 +1155,192 @@ class DeployGeneratorDialog(QDialog):
             return []
         
         return [[combined_results]]
+    
+    def generate_vpes_from_references(self):
+        """
+        Calculates VPEs from the currently highlighted reference images and
+        stores them in self.reference_vpes, overwriting any previous ones.
+        """
+        if not self.loaded_model:
+            QMessageBox.warning(self, "No Model Loaded", "A model must be loaded before generating VPEs.")
+            return
+
+        # Always sync with the live UI selection before generating.
+        self.update_stashed_references_from_ui()
+        references_dict = self._get_references()
+
+        if not references_dict:
+            QMessageBox.information(
+                self,
+                "No References Selected",
+                "Please highlight one or more reference images in the table to generate VPEs."
+            )
+            return
+
+        QApplication.setOverrideCursor(Qt.WaitCursor)
+        progress_bar = ProgressBar(self, title="Generating VPEs")
+        progress_bar.show()
+
+        try:
+            # Make progress bar busy
+            progress_bar.set_busy_mode("Generating VPEs...")
+            # Reload the model to ensure a clean state for VPE generation
+            self.reload_model()
+
+            # The references_to_vpe method will calculate and update self.reference_vpes
+            new_vpes = self.references_to_vpe(references_dict, update_reference_vpes=True)
+
+            if new_vpes:
+                num_vpes = len(new_vpes)
+                num_images = len(references_dict)
+                message = f"Successfully generated {num_vpes} VPEs from {num_images} reference image(s)."
+                self.status_bar.setText(message)
+                QMessageBox.information(self, "VPEs Generated", message)
+            else:
+                message = "Could not generate VPEs. Ensure annotations are valid."
+                self.status_bar.setText(message)
+                QMessageBox.warning(self, "Generation Failed", message)
+
+        except Exception as e:
+            QMessageBox.critical(self, "Error Generating VPEs", f"An unexpected error occurred: {str(e)}")
+            self.status_bar.setText("Error during VPE generation.")
+        finally:
+            QApplication.restoreOverrideCursor()
+            progress_bar.stop_progress()
+            progress_bar.close()
+    
+    def references_to_vpe(self, reference_dict, update_reference_vpes=True):
+        """
+        Converts the contents of a reference dictionary to VPEs (Visual Prompt Embeddings).
+        Reference dictionaries contain information about the visual prompts for each reference image:
+        dict[image_path]: {bboxes, masks, cls}
+
+        Args:
+            reference_dict (dict): The reference dictionary containing visual prompts for each image.
+            update_reference_vpes (bool): Whether to update self.reference_vpes with the results.
+
+        Returns:
+            list: List of individual VPE tensors (normalized), or None if empty reference_dict
+        """
+        # Check if the reference dictionary is empty
+        if not reference_dict:
+            return None
+            
+        # Create a list to hold the individual VPE tensors
+        vpe_list = []
+
+        for ref_path, ref_annotations in reference_dict.items():
+            # Set the prompts to the model predictor
+            self.loaded_model.predictor.set_prompts(ref_annotations)
+
+            # Get the VPE from the model
+            vpe = self.loaded_model.predictor.get_vpe(ref_path)
+            
+            # Normalize individual VPE
+            vpe_normalized = torch.nn.functional.normalize(vpe, p=2, dim=-1)
+            vpe_list.append(vpe_normalized)
+
+        # Check if we have any valid VPEs
+        if not vpe_list:
+            return None
+        
+        # Update the reference_vpes list if requested
+        if update_reference_vpes:
+            self.reference_vpes = vpe_list
+            
+        return vpe_list
+
+    def _apply_model_using_vpe(self, inputs):
+        """
+        Apply the model to the inputs using pre-calculated VPEs from imported files
+        and/or generated from reference annotations.
+        
+        Args:
+            inputs (list): List of input images.
+            
+        Returns:
+            list: List of prediction results.
+        """
+        # First reload the model to clear any cached data
+        self.reload_model()
+        
+        # Initialize combined_vpes list
+        combined_vpes = []
+        
+        # Add imported VPEs if available
+        if self.imported_vpes:
+            combined_vpes.extend(self.imported_vpes)
+            
+        # Add pre-generated reference VPEs if available
+        if self.reference_vpes:
+            combined_vpes.extend(self.reference_vpes)
+        
+        # Check if we have any VPEs to use
+        if not combined_vpes:
+            QMessageBox.warning(
+                self,
+                "No VPEs Available",
+                "No VPEs are available for prediction. "
+                "Please either load a VPE file or generate VPEs from reference images."
+            )
+            return []
+        
+        # Average all the VPEs together to create a final VPE tensor
+        averaged_vpe = torch.cat(combined_vpes).mean(dim=0, keepdim=True)
+        final_vpe = torch.nn.functional.normalize(averaged_vpe, p=2, dim=-1)
+        
+        # For backward compatibility, update self.vpe
+        self.vpe = final_vpe
+        
+        # Set the final VPE to the model
+        self.loaded_model.is_fused = lambda: False 
+        self.loaded_model.set_classes(["object0"], final_vpe)
+        
+        # Make predictions on the target using the averaged VPE
+        results = self.loaded_model.predict(inputs[0],
+                                            visual_prompts=[],
+                                            imgsz=self.imgsz_spinbox.value(),
+                                            conf=self.main_window.get_uncertainty_thresh(),
+                                            iou=self.main_window.get_iou_thresh(),
+                                            max_det=self.get_max_detections(),
+                                            retina_masks=self.task == "segment")
+
+        return [results]
+
+    def _apply_model(self, inputs):
+        """
+        Apply the model to the target inputs. This method handles both image-based 
+        references and VPE-based references.
+        """        
+        # Update the model with user parameters
+        self.task = self.use_task_dropdown.currentText()
+        
+        self.loaded_model.conf = self.main_window.get_uncertainty_thresh()
+        self.loaded_model.iou = self.main_window.get_iou_thresh()
+        self.loaded_model.max_det = self.get_max_detections()
+        
+        # Get the reference information for the currently selected rows
+        references_dict = self._get_references()
+        
+        # Check if the user is using VPE or Reference Images
+        if self.reference_method_combo_box.currentText() == "VPE":
+            # The VPE method will use pre-loaded/pre-generated VPEs.
+            # The internal checks for whether any VPEs exist are now inside _apply_model_using_vpe.
+            results = self._apply_model_using_vpe(inputs)
+        else:  
+            # Use Reference Images method - requires reference images
+            if not references_dict:
+                QMessageBox.warning(
+                    self,
+                    "No References Selected",
+                    "No reference images with valid annotations were selected. "
+                    "Please select at least one reference image."
+                )
+                return []
+                
+            results = self._apply_model_using_images(inputs, references_dict)
+
+        return results
 
     def _apply_sam(self, results_list, image_path):
         """Apply SAM to the results if needed."""
@@ -1000,17 +1437,223 @@ class DeployGeneratorDialog(QDialog):
         progress_bar.stop_progress()
         progress_bar.close()
         
+    def show_vpe(self):
+        """
+        Show a visualization of the currently stored VPEs using PyQtGraph.
+        """
+        try:
+            vpes_with_source = []
+
+            # 1. Add any VPEs that were loaded from a file
+            if self.imported_vpes:
+                for vpe in self.imported_vpes:
+                    vpes_with_source.append((vpe, "Import"))
+
+            # 2. Add any pre-generated VPEs from reference images
+            if self.reference_vpes:
+                for vpe in self.reference_vpes:
+                    vpes_with_source.append((vpe, "Reference"))
+
+            # 3. Check if there is anything to visualize
+            if not vpes_with_source:
+                QMessageBox.warning(
+                    self,
+                    "No VPEs Available",
+                    "No VPEs available to visualize. Please load a VPE file or generate VPEs from references first."
+                )
+                return
+
+            # 4. Create the visualization dialog
+            all_vpe_tensors = [vpe for vpe, source in vpes_with_source]
+            averaged_vpe = torch.cat(all_vpe_tensors).mean(dim=0, keepdim=True)
+            final_vpe = torch.nn.functional.normalize(averaged_vpe, p=2, dim=-1)
+
+            QApplication.setOverrideCursor(Qt.WaitCursor)
+            
+            dialog = VPEVisualizationDialog(vpes_with_source, final_vpe, self)
+            QApplication.restoreOverrideCursor()
+            
+            dialog.exec_()
+            
+        except Exception as e:
+            QApplication.restoreOverrideCursor()
+            QMessageBox.critical(self, "Error Visualizing VPE", f"An error occurred: {str(e)}")
+        
     def deactivate_model(self):
         """
         Deactivate the currently loaded model and clean up resources.
         """
         self.loaded_model = None
         self.model_path = None
-        # Clean up resources
+        
+        # Clear all VPE-related data
+        self.vpe_path_edit.clear()
+        self.vpe_path = None
+        self.vpe = None
+        self.imported_vpes = []
+        self.reference_vpes = []
+        
+        # Clean up references
         gc.collect()
         torch.cuda.empty_cache()
+        
         # Untoggle all tools
         self.main_window.untoggle_all_tools()
+        
         # Update status bar
         self.status_bar.setText("No model loaded")
-        QMessageBox.information(self, "Model Deactivated", "Model deactivated")
+        QMessageBox.information(self, "Model Deactivated", "Model deactivated")
+
+        
+class VPEVisualizationDialog(QDialog):
+    """
+    Dialog for visualizing VPE embeddings in 2D space using PCA.
+    """
+    def __init__(self, vpe_list_with_source, final_vpe=None, parent=None):
+        """
+        Initialize the dialog with a list of VPE tensors and their sources.
+        
+        Args:
+            vpe_list_with_source (list): List of (VPE tensor, source_str) tuples
+            final_vpe (torch.Tensor, optional): The final (averaged) VPE
+            parent (QWidget, optional): Parent widget
+        """
+        super().__init__(parent)
+        self.setWindowTitle("VPE Visualization")
+        self.resize(1000, 1000)
+        
+        # Add a maximize button to the dialog's title bar
+        self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint)
+        
+        # Store the VPEs and their sources
+        self.vpe_list_with_source = vpe_list_with_source
+        self.final_vpe = final_vpe
+        
+        # Create the layout
+        layout = QVBoxLayout(self)
+        
+        # Create the plot widget
+        self.plot_widget = pg.PlotWidget()
+        self.plot_widget.setBackground('w')  # White background
+        self.plot_widget.setTitle("PCA Visualization of Visual Prompt Embeddings", color="#000000", size="10pt")
+        self.plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        
+        # Add the plot widget to the layout
+        layout.addWidget(self.plot_widget)
+        
+        # Add spacing between plot_widget and info_label
+        layout.addSpacing(20)
+        
+        # Add information label at the bottom
+        self.info_label = QLabel()
+        self.info_label.setAlignment(Qt.AlignCenter)
+        layout.addWidget(self.info_label)
+        
+        # Create the button box
+        button_box = QDialogButtonBox(QDialogButtonBox.Close)
+        button_box.rejected.connect(self.reject)
+        layout.addWidget(button_box)
+        
+        # Visualize the VPEs
+        self.visualize_vpes()
+    
+    def visualize_vpes(self):
+        """
+        Apply PCA to the VPE tensors and visualize them in 2D space.
+        """
+        if not self.vpe_list_with_source:
+            self.info_label.setText("No VPEs available to visualize.")
+            return
+        
+        # Convert tensors to numpy arrays for PCA, separating them from the source string
+        vpe_arrays = [vpe.detach().cpu().numpy().squeeze() for vpe, source in self.vpe_list_with_source]
+        
+        # If final VPE is provided, add it to the arrays
+        final_vpe_array = None
+        if self.final_vpe is not None:
+            final_vpe_array = self.final_vpe.detach().cpu().numpy().squeeze()
+            all_vpes = np.vstack(vpe_arrays + [final_vpe_array])
+        else:
+            all_vpes = np.vstack(vpe_arrays)
+        
+        # Apply PCA to reduce to 2 dimensions
+        pca = PCA(n_components=2)
+        vpes_2d = pca.fit_transform(all_vpes)
+        
+        # Clear the plot
+        self.plot_widget.clear()
+        
+        # Generate random colors for individual VPEs
+        num_vpes = len(vpe_arrays)
+        colors = self.generate_distinct_colors(num_vpes)
+        
+        # Create a legend with 3 columns to keep it compact
+        legend = self.plot_widget.addLegend(colCount=3)
+        
+        # Plot individual VPEs
+        for i, (vpe_tuple, vpe_2d) in enumerate(zip(self.vpe_list_with_source, vpes_2d[:num_vpes])):
+            source_char = 'I' if vpe_tuple[1] == 'Import' else 'R'
+            color = pg.mkColor(colors[i])
+            scatter = pg.ScatterPlotItem(
+                x=[vpe_2d[0]], 
+                y=[vpe_2d[1]], 
+                brush=color, 
+                size=15,
+                name=f"VPE {i+1} ({source_char})"
+            )
+            self.plot_widget.addItem(scatter)
+        
+        # Plot the final (averaged) VPE if available
+        if final_vpe_array is not None:
+            final_vpe_2d = vpes_2d[-1]
+            scatter = pg.ScatterPlotItem(
+                x=[final_vpe_2d[0]], 
+                y=[final_vpe_2d[1]], 
+                brush=pg.mkBrush(color='r'), 
+                size=20,
+                symbol='star',
+                name="Final VPE"
+            )
+            self.plot_widget.addItem(scatter)
+        
+        # Update the information label
+        orig_dim = self.vpe_list_with_source[0][0].shape[-1]
+        explained_variance = sum(pca.explained_variance_ratio_)
+        self.info_label.setText(
+            f"Original dimension: {orig_dim} → Reduced to 2D\n"
+            f"Total explained variance: {explained_variance:.2%}\n"
+            f"PC1: {pca.explained_variance_ratio_[0]:.2%} variance, "
+            f"PC2: {pca.explained_variance_ratio_[1]:.2%} variance"
+        )
+    
+    def generate_distinct_colors(self, num_colors):
+        """
+        Generate visually distinct colors by using evenly spaced hues
+        with random saturation and value.
+        
+        Args:
+            num_colors (int): Number of colors to generate
+            
+        Returns:
+            list: List of color hex strings
+        """
+        import random
+        from colorsys import hsv_to_rgb
+        
+        colors = []
+        for i in range(num_colors):
+            # Use golden ratio to space hues evenly
+            hue = (i * 0.618033988749895) % 1.0
+            # Random saturation between 0.6-1.0 (avoid too pale)
+            saturation = random.uniform(0.6, 1.0)
+            # Random value between 0.7-1.0 (avoid too dark)
+            value = random.uniform(0.7, 1.0)
+            
+            # Convert HSV to RGB (0-1 range)
+            r, g, b = hsv_to_rgb(hue, saturation, value)
+            
+            # Convert RGB to hex string
+            hex_color = f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}"
+            colors.append(hex_color)
+        
+        return colors