""" Annotation tab for the microscopy object detection application. Manual annotation with pen tool and object class management. """ from PySide6.QtWidgets import ( QWidget, QVBoxLayout, QHBoxLayout, QLabel, QGroupBox, QPushButton, QFileDialog, QMessageBox, QSplitter, ) from PySide6.QtCore import Qt, QSettings from pathlib import Path from src.database.db_manager import DatabaseManager from src.utils.config_manager import ConfigManager from src.utils.image import Image, ImageLoadError from src.utils.logger import get_logger from src.gui.widgets import AnnotationCanvasWidget, AnnotationToolsWidget logger = get_logger(__name__) class AnnotationTab(QWidget): """Annotation tab for manual image annotation.""" def __init__( self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None ): super().__init__(parent) self.db_manager = db_manager self.config_manager = config_manager self.current_image = None self.current_image_path = None self.current_image_id = None self.current_annotations = [] self._setup_ui() def _setup_ui(self): """Setup user interface.""" layout = QVBoxLayout() # Main horizontal splitter to divide left (image) and right (controls) self.main_splitter = QSplitter(Qt.Horizontal) self.main_splitter.setHandleWidth(10) # { Left splitter for image display and zoom info self.left_splitter = QSplitter(Qt.Vertical) self.left_splitter.setHandleWidth(10) # Annotation canvas section canvas_group = QGroupBox("Annotation Canvas") canvas_layout = QVBoxLayout() # Use the AnnotationCanvasWidget self.annotation_canvas = AnnotationCanvasWidget() self.annotation_canvas.zoom_changed.connect(self._on_zoom_changed) self.annotation_canvas.annotation_drawn.connect(self._on_annotation_drawn) canvas_layout.addWidget(self.annotation_canvas) canvas_group.setLayout(canvas_layout) self.left_splitter.addWidget(canvas_group) # Controls info controls_info = QLabel( "Zoom: Mouse wheel or +/- keys | Drawing: Enable pen and drag mouse" ) controls_info.setStyleSheet("QLabel { color: #888; font-style: italic; }") self.left_splitter.addWidget(controls_info) # } # { Right splitter for annotation tools and controls self.right_splitter = QSplitter(Qt.Vertical) self.right_splitter.setHandleWidth(10) # Annotation tools section self.annotation_tools = AnnotationToolsWidget(self.db_manager) self.annotation_tools.polyline_enabled_changed.connect( self.annotation_canvas.set_polyline_enabled ) self.annotation_tools.polyline_pen_color_changed.connect( self.annotation_canvas.set_polyline_pen_color ) self.annotation_tools.polyline_pen_width_changed.connect( self.annotation_canvas.set_polyline_pen_width ) # Show / hide bounding boxes self.annotation_tools.show_bboxes_changed.connect( self.annotation_canvas.set_show_bboxes ) # RDP simplification controls self.annotation_tools.simplify_on_finish_changed.connect( self._on_simplify_on_finish_changed ) self.annotation_tools.simplify_epsilon_changed.connect( self._on_simplify_epsilon_changed ) # Class selection and class-color changes self.annotation_tools.class_selected.connect(self._on_class_selected) self.annotation_tools.class_color_changed.connect(self._on_class_color_changed) self.annotation_tools.clear_annotations_requested.connect( self._on_clear_annotations ) self.right_splitter.addWidget(self.annotation_tools) # Image loading section load_group = QGroupBox("Image Loading") load_layout = QVBoxLayout() # Load image button button_layout = QHBoxLayout() self.load_image_btn = QPushButton("Load Image") self.load_image_btn.clicked.connect(self._load_image) button_layout.addWidget(self.load_image_btn) button_layout.addStretch() load_layout.addLayout(button_layout) # Image info label self.image_info_label = QLabel("No image loaded") load_layout.addWidget(self.image_info_label) load_group.setLayout(load_layout) self.right_splitter.addWidget(load_group) # } # Add both splitters to the main horizontal splitter self.main_splitter.addWidget(self.left_splitter) self.main_splitter.addWidget(self.right_splitter) # Set initial sizes: 75% for left (image), 25% for right (controls) self.main_splitter.setSizes([750, 250]) layout.addWidget(self.main_splitter) self.setLayout(layout) # Restore splitter positions from settings self._restore_state() def _load_image(self): """Load and display an image file.""" # Get last opened directory from QSettings settings = QSettings("microscopy_app", "object_detection") last_dir = settings.value("annotation_tab/last_directory", None) # Fallback to image repository path or home directory if last_dir and Path(last_dir).exists(): start_dir = last_dir else: repo_path = self.config_manager.get_image_repository_path() start_dir = repo_path if repo_path else str(Path.home()) # Open file dialog file_path, _ = QFileDialog.getOpenFileName( self, "Select Image", start_dir, "Images (*.jpg *.jpeg *.png *.tif *.tiff *.bmp)", ) if not file_path: return try: # Load image using Image class self.current_image = Image(file_path) self.current_image_path = file_path # Store the directory for next time settings.setValue( "annotation_tab/last_directory", str(Path(file_path).parent) ) # Get or create image in database relative_path = str(Path(file_path).name) # Simplified for now self.current_image_id = self.db_manager.get_or_create_image( relative_path, Path(file_path).name, self.current_image.width, self.current_image.height, ) # Display image using the AnnotationCanvasWidget self.annotation_canvas.load_image(self.current_image) # Load and display any existing annotations for this image self._load_annotations_for_current_image() # Update info label self._update_image_info() logger.info(f"Loaded image: {file_path} (DB ID: {self.current_image_id})") except ImageLoadError as e: logger.error(f"Failed to load image: {e}") QMessageBox.critical( self, "Error Loading Image", f"Failed to load image:\n{str(e)}" ) except Exception as e: logger.error(f"Unexpected error loading image: {e}") QMessageBox.critical(self, "Error", f"Unexpected error:\n{str(e)}") def _update_image_info(self): """Update the image info label with current image details.""" if self.current_image is None: self.image_info_label.setText("No image loaded") return zoom_percentage = self.annotation_canvas.get_zoom_percentage() info_text = ( f"File: {Path(self.current_image_path).name}\n" f"Size: {self.current_image.width}x{self.current_image.height} pixels\n" f"Channels: {self.current_image.channels}\n" f"Data type: {self.current_image.dtype}\n" f"Format: {self.current_image.format.upper()}\n" f"File size: {self.current_image.size_mb:.2f} MB\n" f"Zoom: {zoom_percentage}%" ) self.image_info_label.setText(info_text) def _on_zoom_changed(self, zoom_scale: float): """Handle zoom level changes from the annotation canvas.""" self._update_image_info() def _on_annotation_drawn(self, points: list): """ Handle when an annotation stroke is drawn. Saves the new annotation directly to the database and refreshes the on-canvas display of annotations for the current image. """ # Ensure we have an image loaded and in the DB if not self.current_image or not self.current_image_id: logger.warning("Annotation drawn but no image loaded") QMessageBox.warning( self, "No Image", "Please load an image before drawing annotations.", ) return current_class = self.annotation_tools.get_current_class() if not current_class: logger.warning("Annotation drawn but no object class selected") QMessageBox.warning( self, "No Class Selected", "Please select an object class before drawing annotations.", ) return if not points: logger.warning("Annotation drawn with no points, ignoring") return # points are [(x_norm, y_norm), ...] xs = [p[0] for p in points] ys = [p[1] for p in points] x_min, x_max = min(xs), max(xs) y_min, y_max = min(ys), max(ys) # Store segmentation mask in [y_norm, x_norm] format to match DB db_polyline = [[float(y), float(x)] for (x, y) in points] try: annotation_id = self.db_manager.add_annotation( image_id=self.current_image_id, class_id=current_class["id"], bbox=(x_min, y_min, x_max, y_max), annotator="manual", segmentation_mask=db_polyline, verified=False, ) logger.info( f"Saved annotation (ID: {annotation_id}) for class " f"'{current_class['class_name']}' " f"Bounding box: ({x_min:.3f}, {y_min:.3f}) to ({x_max:.3f}, {y_max:.3f})\n" f"with {len(points)} polyline points" ) # Reload annotations from DB and redraw (respecting current class filter) self._load_annotations_for_current_image() except Exception as e: logger.error(f"Failed to save annotation: {e}") QMessageBox.critical(self, "Error", f"Failed to save annotation:\n{str(e)}") def _on_simplify_on_finish_changed(self, enabled: bool): """Update canvas simplify-on-finish flag from tools widget.""" self.annotation_canvas.simplify_on_finish = enabled logger.debug(f"Annotation simplification on finish set to {enabled}") def _on_simplify_epsilon_changed(self, epsilon: float): """Update canvas RDP epsilon from tools widget.""" self.annotation_canvas.simplify_epsilon = float(epsilon) logger.debug(f"Annotation simplification epsilon set to {epsilon}") def _on_class_color_changed(self): """ Handle changes to the selected object's class color. When the user updates a class color in the tools widget, reload the annotations for the current image so that all polylines are redrawn using the updated per-class colors. """ if not self.current_image_id: return logger.debug( f"Class color changed; reloading annotations for image ID {self.current_image_id}" ) self._load_annotations_for_current_image() def _on_class_selected(self, class_data): """ Handle when an object class is selected or cleared. When a specific class is selected, only annotations of that class are drawn. When the selection is cleared (\"-- Select Class --\"), all annotations are shown. """ if class_data: logger.debug(f"Object class selected: {class_data['class_name']}") else: logger.debug( 'No class selected ("-- Select Class --"), showing all annotations' ) # Whenever the selection changes, update which annotations are visible self._redraw_annotations_for_current_filter() def _on_clear_annotations(self): """Handle clearing all annotations.""" self.annotation_canvas.clear_annotations() logger.info("Cleared all annotations") def _load_annotations_for_current_image(self): """ Load all annotations for the current image from the database and redraw them on the canvas, honoring the currently selected class filter (if any). """ if not self.current_image_id: self.current_annotations = [] self.annotation_canvas.clear_annotations() return try: self.current_annotations = self.db_manager.get_annotations_for_image( self.current_image_id ) self._redraw_annotations_for_current_filter() except Exception as e: logger.error( f"Failed to load annotations for image {self.current_image_id}: {e}" ) QMessageBox.critical( self, "Error", f"Failed to load annotations for this image:\n{str(e)}", ) def _redraw_annotations_for_current_filter(self): """ Redraw annotations for the current image, optionally filtered by the currently selected object class. """ # Clear current on-canvas annotations but keep the image self.annotation_canvas.clear_annotations() if not self.current_annotations: return current_class = self.annotation_tools.get_current_class() selected_class_id = current_class["id"] if current_class else None drawn_count = 0 for ann in self.current_annotations: # Filter by class if one is selected if ( selected_class_id is not None and ann.get("class_id") != selected_class_id ): continue if ann.get("segmentation_mask"): polyline = ann["segmentation_mask"] color = ann.get("class_color", "#FF0000") self.annotation_canvas.draw_saved_polyline(polyline, color, width=3) self.annotation_canvas.draw_saved_bbox( [ann["x_min"], ann["y_min"], ann["x_max"], ann["y_max"]], color, width=3, ) drawn_count += 1 logger.info( f"Displayed {drawn_count} annotation(s) for current image with " f"{'no class filter' if selected_class_id is None else f'class_id={selected_class_id}'}" ) def _restore_state(self): """Restore splitter positions from settings.""" settings = QSettings("microscopy_app", "object_detection") # Restore main splitter state main_state = settings.value("annotation_tab/main_splitter_state") if main_state: self.main_splitter.restoreState(main_state) logger.debug("Restored main splitter state") # Restore left splitter state left_state = settings.value("annotation_tab/left_splitter_state") if left_state: self.left_splitter.restoreState(left_state) logger.debug("Restored left splitter state") # Restore right splitter state right_state = settings.value("annotation_tab/right_splitter_state") if right_state: self.right_splitter.restoreState(right_state) logger.debug("Restored right splitter state") def save_state(self): """Save splitter positions to settings.""" settings = QSettings("microscopy_app", "object_detection") # Save main splitter state settings.setValue( "annotation_tab/main_splitter_state", self.main_splitter.saveState() ) # Save left splitter state settings.setValue( "annotation_tab/left_splitter_state", self.left_splitter.saveState() ) # Save right splitter state settings.setValue( "annotation_tab/right_splitter_state", self.right_splitter.saveState() ) logger.debug("Saved annotation tab splitter states") def refresh(self): """Refresh the tab.""" pass