""" Annotation tools widget for controlling annotation parameters. Includes polyline tool, color picker, class selection, and annotation management. """ from PySide6.QtWidgets import ( QWidget, QVBoxLayout, QHBoxLayout, QLabel, QGroupBox, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox, QCheckBox, QColorDialog, QInputDialog, QMessageBox, ) from PySide6.QtGui import QColor, QIcon, QPixmap, QPainter from PySide6.QtCore import Qt, Signal from typing import Optional, Dict from src.database.db_manager import DatabaseManager from src.utils.logger import get_logger logger = get_logger(__name__) class AnnotationToolsWidget(QWidget): """ Widget for annotation tool controls. Features: - Enable/disable polyline tool - Color selection for polyline pen - Object class selection - Add new object classes - Pen width control - Clear annotations Signals: polyline_enabled_changed: Emitted when polyline tool is enabled/disabled (bool) polyline_pen_color_changed: Emitted when polyline pen color changes (QColor) polyline_pen_width_changed: Emitted when polyline pen width changes (int) class_selected: Emitted when object class is selected (dict) clear_annotations_requested: Emitted when clear button is pressed """ polyline_enabled_changed = Signal(bool) polyline_pen_color_changed = Signal(QColor) polyline_pen_width_changed = Signal(int) simplify_on_finish_changed = Signal(bool) simplify_epsilon_changed = Signal(float) # Toggle visibility of bounding boxes on the canvas show_bboxes_changed = Signal(bool) class_selected = Signal(dict) class_color_changed = Signal() clear_annotations_requested = Signal() def __init__(self, db_manager: DatabaseManager, parent=None): """ Initialize annotation tools widget. Args: db_manager: Database manager instance parent: Parent widget """ super().__init__(parent) self.db_manager = db_manager self.polyline_enabled = False self.current_color = QColor(255, 0, 0, 128) # Red with 50% alpha self.current_class = None self._setup_ui() self._load_object_classes() def _setup_ui(self): """Setup user interface.""" layout = QVBoxLayout() # Polyline Tool Group polyline_group = QGroupBox("Polyline Tool") polyline_layout = QVBoxLayout() # Enable/Disable polyline tool button_layout = QHBoxLayout() self.polyline_toggle_btn = QPushButton("Start Drawing Polyline") self.polyline_toggle_btn.setCheckable(True) self.polyline_toggle_btn.clicked.connect(self._on_polyline_toggle) button_layout.addWidget(self.polyline_toggle_btn) polyline_layout.addLayout(button_layout) # Polyline pen width control width_layout = QHBoxLayout() width_layout.addWidget(QLabel("Pen Width:")) self.polyline_pen_width_spin = QSpinBox() self.polyline_pen_width_spin.setMinimum(1) self.polyline_pen_width_spin.setMaximum(20) self.polyline_pen_width_spin.setValue(3) self.polyline_pen_width_spin.valueChanged.connect( self._on_polyline_pen_width_changed ) width_layout.addWidget(self.polyline_pen_width_spin) width_layout.addStretch() polyline_layout.addLayout(width_layout) # Simplification controls (RDP) simplify_layout = QHBoxLayout() self.simplify_checkbox = QCheckBox("Simplify on finish") self.simplify_checkbox.setChecked(True) self.simplify_checkbox.stateChanged.connect(self._on_simplify_toggle) simplify_layout.addWidget(self.simplify_checkbox) simplify_layout.addWidget(QLabel("epsilon (px):")) self.eps_spin = QDoubleSpinBox() self.eps_spin.setRange(0.0, 1000.0) self.eps_spin.setSingleStep(0.5) self.eps_spin.setValue(2.0) self.eps_spin.valueChanged.connect(self._on_eps_change) simplify_layout.addWidget(self.eps_spin) simplify_layout.addStretch() polyline_layout.addLayout(simplify_layout) polyline_group.setLayout(polyline_layout) layout.addWidget(polyline_group) # Object Class Group class_group = QGroupBox("Object Class") class_layout = QVBoxLayout() # Class selection dropdown self.class_combo = QComboBox() self.class_combo.currentIndexChanged.connect(self._on_class_selected) class_layout.addWidget(self.class_combo) # Add / manage classes class_button_layout = QHBoxLayout() self.add_class_btn = QPushButton("Add New Class") self.add_class_btn.clicked.connect(self._on_add_class) class_button_layout.addWidget(self.add_class_btn) self.refresh_classes_btn = QPushButton("Refresh") self.refresh_classes_btn.clicked.connect(self._load_object_classes) class_button_layout.addWidget(self.refresh_classes_btn) class_layout.addLayout(class_button_layout) # Class color (associated with selected object class) color_layout = QHBoxLayout() color_layout.addWidget(QLabel("Class Color:")) self.color_btn = QPushButton() self.color_btn.setFixedSize(40, 30) self.color_btn.clicked.connect(self._on_color_picker) self._update_color_button() color_layout.addWidget(self.color_btn) color_layout.addStretch() class_layout.addLayout(color_layout) # Selected class info self.class_info_label = QLabel("No class selected") self.class_info_label.setWordWrap(True) self.class_info_label.setStyleSheet( "QLabel { color: #888; font-style: italic; }" ) class_layout.addWidget(self.class_info_label) class_group.setLayout(class_layout) layout.addWidget(class_group) # Actions Group actions_group = QGroupBox("Actions") actions_layout = QVBoxLayout() # Show / hide bounding boxes self.show_bboxes_checkbox = QCheckBox("Show bounding boxes") self.show_bboxes_checkbox.setChecked(True) self.show_bboxes_checkbox.stateChanged.connect(self._on_show_bboxes_toggle) actions_layout.addWidget(self.show_bboxes_checkbox) self.clear_btn = QPushButton("Clear All Annotations") self.clear_btn.clicked.connect(self._on_clear_annotations) actions_layout.addWidget(self.clear_btn) actions_group.setLayout(actions_layout) layout.addWidget(actions_group) layout.addStretch() self.setLayout(layout) def _update_color_button(self): """Update the color button appearance with current color.""" pixmap = QPixmap(40, 30) pixmap.fill(self.current_color) # Add border painter = QPainter(pixmap) painter.setPen(Qt.black) painter.drawRect(0, 0, pixmap.width() - 1, pixmap.height() - 1) painter.end() self.color_btn.setIcon(QIcon(pixmap)) self.color_btn.setStyleSheet(f"background-color: {self.current_color.name()};") def _load_object_classes(self): """Load object classes from database and populate combo box.""" try: classes = self.db_manager.get_object_classes() # Clear and repopulate combo box self.class_combo.clear() self.class_combo.addItem("-- Select Class / Show All --", None) for cls in classes: self.class_combo.addItem(cls["class_name"], cls) logger.debug(f"Loaded {len(classes)} object classes") except Exception as e: logger.error(f"Error loading object classes: {e}") QMessageBox.warning( self, "Error", f"Failed to load object classes:\n{str(e)}" ) def _on_polyline_toggle(self, checked: bool): """Handle polyline tool enable/disable.""" self.polyline_enabled = checked if checked: self.polyline_toggle_btn.setText("Stop Drawing Polyline") self.polyline_toggle_btn.setStyleSheet( "QPushButton { background-color: #4CAF50; }" ) else: self.polyline_toggle_btn.setText("Start Drawing Polyline") self.polyline_toggle_btn.setStyleSheet("") self.polyline_enabled_changed.emit(self.polyline_enabled) logger.debug(f"Polyline tool {'enabled' if checked else 'disabled'}") def _on_polyline_pen_width_changed(self, width: int): """Handle polyline pen width changes.""" self.polyline_pen_width_changed.emit(width) logger.debug(f"Polyline pen width changed to {width}") def _on_simplify_toggle(self, state: int): """Handle simplify-on-finish checkbox toggle.""" enabled = bool(state) self.simplify_on_finish_changed.emit(enabled) logger.debug(f"Simplify on finish set to {enabled}") def _on_eps_change(self, val: float): """Handle epsilon (RDP tolerance) value changes.""" epsilon = float(val) self.simplify_epsilon_changed.emit(epsilon) logger.debug(f"Simplification epsilon changed to {epsilon}") def _on_show_bboxes_toggle(self, state: int): """Handle 'Show bounding boxes' checkbox toggle.""" show = bool(state) self.show_bboxes_changed.emit(show) logger.debug(f"Show bounding boxes set to {show}") def _on_color_picker(self): """Open color picker dialog and update the selected object's class color.""" if not self.current_class: QMessageBox.warning( self, "No Class Selected", "Please select an object class before changing its color.", ) return # Use current class color (without alpha) as the base base_color = QColor(self.current_class.get("color", self.current_color.name())) color = QColorDialog.getColor( base_color, self, "Select Class Color", QColorDialog.ShowAlphaChannel, # Allow alpha in UI, but store RGB in DB ) if not color.isValid(): return # Normalize to opaque RGB for storage new_color = QColor(color) new_color.setAlpha(255) hex_color = new_color.name() try: # Update in database self.db_manager.update_object_class( class_id=self.current_class["id"], color=hex_color ) except Exception as e: logger.error(f"Failed to update class color in database: {e}") QMessageBox.critical( self, "Error", f"Failed to update class color in database:\n{str(e)}", ) return # Update local class data and combo box item data self.current_class["color"] = hex_color current_index = self.class_combo.currentIndex() if current_index >= 0: self.class_combo.setItemData(current_index, dict(self.current_class)) # Update info label text info_text = f"Class: {self.current_class['class_name']}\nColor: {hex_color}" if self.current_class.get("description"): info_text += f"\nDescription: {self.current_class['description']}" self.class_info_label.setText(info_text) # Use semi-transparent version for polyline pen / button preview class_color = QColor(hex_color) class_color.setAlpha(128) self.current_color = class_color self._update_color_button() self.polyline_pen_color_changed.emit(class_color) logger.debug( f"Updated class '{self.current_class['class_name']}' color to " f"{hex_color} (polyline pen alpha={class_color.alpha()})" ) # Notify listeners (e.g., AnnotationTab) so they can reload/redraw self.class_color_changed.emit() def _on_class_selected(self, index: int): """Handle object class selection (including '-- Select Class --').""" class_data = self.class_combo.currentData() if class_data: self.current_class = class_data # Update info label info_text = ( f"Class: {class_data['class_name']}\n" f"Color: {class_data['color']}" ) if class_data.get("description"): info_text += f"\nDescription: {class_data['description']}" self.class_info_label.setText(info_text) # Update polyline pen color to match class color with semi-transparency class_color = QColor(class_data["color"]) if class_color.isValid(): # Add 50% alpha for semi-transparency class_color.setAlpha(128) self.current_color = class_color self._update_color_button() self.polyline_pen_color_changed.emit(class_color) self.class_selected.emit(class_data) logger.debug(f"Selected class: {class_data['class_name']}") else: # "-- Select Class --" chosen: clear current class and show all annotations self.current_class = None self.class_info_label.setText("No class selected") self.class_selected.emit(None) logger.debug("Class selection cleared: showing annotations for all classes") def _on_add_class(self): """Handle adding a new object class.""" # Get class name class_name, ok = QInputDialog.getText( self, "Add Object Class", "Enter class name:" ) if not ok or not class_name.strip(): return class_name = class_name.strip() # Check if class already exists existing = self.db_manager.get_object_class_by_name(class_name) if existing: QMessageBox.warning( self, "Class Exists", f"A class named '{class_name}' already exists." ) return # Get color color = QColorDialog.getColor(self.current_color, self, "Select Class Color") if not color.isValid(): return # Get optional description description, ok = QInputDialog.getText( self, "Class Description", "Enter class description (optional):" ) if not ok: description = None # Add to database try: class_id = self.db_manager.add_object_class( class_name, color.name(), description.strip() if description else None ) logger.info(f"Added new object class: {class_name} (ID: {class_id})") # Reload classes and select the new one self._load_object_classes() # Find and select the newly added class for i in range(self.class_combo.count()): class_data = self.class_combo.itemData(i) if class_data and class_data.get("id") == class_id: self.class_combo.setCurrentIndex(i) break QMessageBox.information( self, "Success", f"Class '{class_name}' added successfully!" ) except Exception as e: logger.error(f"Error adding object class: {e}") QMessageBox.critical( self, "Error", f"Failed to add object class:\n{str(e)}" ) def _on_clear_annotations(self): """Handle clear annotations button.""" reply = QMessageBox.question( self, "Clear Annotations", "Are you sure you want to clear all annotations?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No, ) if reply == QMessageBox.Yes: self.clear_annotations_requested.emit() logger.debug("Clear annotations requested") def get_current_class(self) -> Optional[Dict]: """Get currently selected object class.""" return self.current_class def get_polyline_pen_color(self) -> QColor: """Get current polyline pen color.""" return self.current_color def get_polyline_pen_width(self) -> int: """Get current polyline pen width.""" return self.polyline_pen_width_spin.value() def is_polyline_enabled(self) -> bool: """Check if polyline tool is enabled.""" return self.polyline_enabled