""" Annotation tools widget for controlling annotation parameters. Includes pen 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 pen tool - Color selection for pen - Object class selection - Add new object classes - Pen width control - Clear annotations Signals: pen_enabled_changed: Emitted when pen tool is enabled/disabled (bool) pen_color_changed: Emitted when pen color changes (QColor) pen_width_changed: Emitted when pen width changes (int) class_selected: Emitted when object class is selected (dict) clear_annotations_requested: Emitted when clear button is pressed """ pen_enabled_changed = Signal(bool) pen_color_changed = Signal(QColor) pen_width_changed = Signal(int) simplify_on_finish_changed = Signal(bool) simplify_epsilon_changed = Signal(float) class_selected = Signal(dict) clear_annotations_requested = Signal() process_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.pen_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() # Pen Tool Group pen_group = QGroupBox("Pen Tool") pen_layout = QVBoxLayout() # Enable/Disable pen button_layout = QHBoxLayout() self.pen_toggle_btn = QPushButton("Enable Pen") self.pen_toggle_btn.setCheckable(True) self.pen_toggle_btn.clicked.connect(self._on_pen_toggle) button_layout.addWidget(self.pen_toggle_btn) pen_layout.addLayout(button_layout) # Pen width control width_layout = QHBoxLayout() width_layout.addWidget(QLabel("Pen Width:")) self.pen_width_spin = QSpinBox() self.pen_width_spin.setMinimum(1) self.pen_width_spin.setMaximum(20) self.pen_width_spin.setValue(3) self.pen_width_spin.valueChanged.connect(self._on_pen_width_changed) width_layout.addWidget(self.pen_width_spin) width_layout.addStretch() pen_layout.addLayout(width_layout) # Color selection color_layout = QHBoxLayout() color_layout.addWidget(QLabel("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() pen_layout.addLayout(color_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() pen_layout.addLayout(simplify_layout) pen_group.setLayout(pen_layout) layout.addWidget(pen_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 class button 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) # 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() self.process_btn = QPushButton("Process Annotations") self.process_btn.clicked.connect(self._on_process_annotations) self.process_btn.setStyleSheet( "QPushButton { background-color: #2196F3; color: white; font-weight: bold; }" ) actions_layout.addWidget(self.process_btn) 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 --", 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_pen_toggle(self, checked: bool): """Handle pen tool enable/disable.""" self.pen_enabled = checked if checked: self.pen_toggle_btn.setText("Disable Pen") self.pen_toggle_btn.setStyleSheet( "QPushButton { background-color: #4CAF50; }" ) else: self.pen_toggle_btn.setText("Enable Pen") self.pen_toggle_btn.setStyleSheet("") self.pen_enabled_changed.emit(self.pen_enabled) logger.debug(f"Pen tool {'enabled' if checked else 'disabled'}") def _on_pen_width_changed(self, width: int): """Handle pen width changes.""" self.pen_width_changed.emit(width) logger.debug(f"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_color_picker(self): """Open color picker dialog with alpha support.""" color = QColorDialog.getColor( self.current_color, self, "Select Pen Color", QColorDialog.ShowAlphaChannel, # Enable alpha channel selection ) if color.isValid(): self.current_color = color self._update_color_button() self.pen_color_changed.emit(color) logger.debug( f"Pen color changed to {color.name()} with alpha {color.alpha()}" ) def _on_class_selected(self, index: int): """Handle object class selection.""" 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 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.pen_color_changed.emit(class_color) self.class_selected.emit(class_data) logger.debug(f"Selected class: {class_data['class_name']}") else: self.current_class = None self.class_info_label.setText("No class selected") 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 _on_process_annotations(self): """Handle process annotations button.""" if not self.current_class: QMessageBox.warning( self, "No Class Selected", "Please select an object class before processing annotations.", ) return self.process_annotations_requested.emit() logger.debug("Process annotations requested") def get_current_class(self) -> Optional[Dict]: """Get currently selected object class.""" return self.current_class def get_pen_color(self) -> QColor: """Get current pen color.""" return self.current_color def get_pen_width(self) -> int: """Get current pen width.""" return self.pen_width_spin.value() def is_pen_enabled(self) -> bool: """Check if pen tool is enabled.""" return self.pen_enabled