7 Commits

8 changed files with 2450 additions and 472 deletions

View File

@@ -18,6 +18,8 @@ training:
default_imgsz: 640
default_patience: 50
default_lr0: 0.01
last_dataset_yaml: /home/martin/code/object_detection/data/datasets/data.yaml
last_dataset_dir: /home/martin/code/object_detection/data/datasets
detection:
default_confidence: 0.25
default_iou: 0.45

View File

@@ -10,6 +10,13 @@ from typing import List, Dict, Optional, Tuple, Any, Union
from pathlib import Path
import csv
import hashlib
import yaml
from src.utils.logger import get_logger
IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp")
logger = get_logger(__name__)
class DatabaseManager:
@@ -706,6 +713,25 @@ class DatabaseManager:
finally:
conn.close()
def delete_annotation(self, annotation_id: int) -> bool:
"""
Delete a manual annotation by ID.
Args:
annotation_id: ID of the annotation to delete
Returns:
True if an annotation was deleted, False otherwise.
"""
conn = self.get_connection()
try:
cursor = conn.cursor()
cursor.execute("DELETE FROM annotations WHERE id = ?", (annotation_id,))
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
# ==================== Object Class Operations ====================
def get_object_classes(self) -> List[Dict]:
@@ -842,6 +868,187 @@ class DatabaseManager:
finally:
conn.close()
# ==================== Dataset Utilities ====================
def compose_data_yaml(
self,
dataset_root: str,
output_path: Optional[str] = None,
splits: Optional[Dict[str, str]] = None,
) -> str:
"""
Compose a YOLO data.yaml file based on dataset folders and database metadata.
Args:
dataset_root: Base directory containing the dataset structure.
output_path: Optional output path; defaults to <dataset_root>/data.yaml.
splits: Optional mapping overriding train/val/test image directories (relative
to dataset_root or absolute paths).
Returns:
Path to the generated YAML file.
"""
dataset_root_path = Path(dataset_root).expanduser()
if not dataset_root_path.exists():
raise ValueError(f"Dataset root does not exist: {dataset_root_path}")
dataset_root_path = dataset_root_path.resolve()
split_map: Dict[str, str] = {key: "" for key in ("train", "val", "test")}
if splits:
for key, value in splits.items():
if key in split_map and value:
split_map[key] = value
inferred = self._infer_split_dirs(dataset_root_path)
for key in split_map:
if not split_map[key]:
split_map[key] = inferred.get(key, "")
for required in ("train", "val"):
if not split_map[required]:
raise ValueError(
"Unable to determine %s image directory under %s. Provide it "
"explicitly via the 'splits' argument."
% (required, dataset_root_path)
)
yaml_splits: Dict[str, str] = {}
for key, value in split_map.items():
if not value:
continue
yaml_splits[key] = self._normalize_split_value(value, dataset_root_path)
class_names = self._fetch_annotation_class_names()
if not class_names:
class_names = [cls["class_name"] for cls in self.get_object_classes()]
if not class_names:
raise ValueError("No object classes available to populate data.yaml")
names_map = {idx: name for idx, name in enumerate(class_names)}
payload: Dict[str, Any] = {
"path": dataset_root_path.as_posix(),
"train": yaml_splits["train"],
"val": yaml_splits["val"],
"names": names_map,
"nc": len(class_names),
}
if yaml_splits.get("test"):
payload["test"] = yaml_splits["test"]
output_path_obj = (
Path(output_path).expanduser()
if output_path
else dataset_root_path / "data.yaml"
)
output_path_obj.parent.mkdir(parents=True, exist_ok=True)
with open(output_path_obj, "w", encoding="utf-8") as handle:
yaml.safe_dump(payload, handle, sort_keys=False)
logger.info(f"Generated data.yaml at {output_path_obj}")
return output_path_obj.as_posix()
def _fetch_annotation_class_names(self) -> List[str]:
"""Return class names referenced by annotations (ordered by class ID)."""
conn = self.get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"""
SELECT DISTINCT c.id, c.class_name
FROM annotations a
JOIN object_classes c ON a.class_id = c.id
ORDER BY c.id
"""
)
rows = cursor.fetchall()
return [row["class_name"] for row in rows]
finally:
conn.close()
def _infer_split_dirs(self, dataset_root: Path) -> Dict[str, str]:
"""Infer train/val/test image directories relative to dataset_root."""
patterns = {
"train": [
"train/images",
"training/images",
"images/train",
"images/training",
"train",
"training",
],
"val": [
"val/images",
"validation/images",
"images/val",
"images/validation",
"val",
"validation",
],
"test": [
"test/images",
"testing/images",
"images/test",
"images/testing",
"test",
"testing",
],
}
inferred: Dict[str, str] = {key: "" for key in patterns}
for split_name, options in patterns.items():
for relative in options:
candidate = (dataset_root / relative).resolve()
if (
candidate.exists()
and candidate.is_dir()
and self._directory_has_images(candidate)
):
try:
inferred[split_name] = candidate.relative_to(
dataset_root
).as_posix()
except ValueError:
inferred[split_name] = candidate.as_posix()
break
return inferred
def _normalize_split_value(self, split_value: str, dataset_root: Path) -> str:
"""Validate and normalize a split directory to a YAML-friendly string."""
split_path = Path(split_value).expanduser()
if not split_path.is_absolute():
split_path = (dataset_root / split_path).resolve()
else:
split_path = split_path.resolve()
if not split_path.exists() or not split_path.is_dir():
raise ValueError(f"Split directory not found: {split_path}")
if not self._directory_has_images(split_path):
raise ValueError(f"No images found under {split_path}")
try:
return split_path.relative_to(dataset_root).as_posix()
except ValueError:
return split_path.as_posix()
@staticmethod
def _directory_has_images(directory: Path, max_checks: int = 2000) -> bool:
"""Return True if directory tree contains at least one image file."""
checked = 0
try:
for file_path in directory.rglob("*"):
if not file_path.is_file():
continue
if file_path.suffix.lower() in IMAGE_EXTENSIONS:
return True
checked += 1
if checked >= max_checks:
break
except Exception:
return False
return False
@staticmethod
def calculate_checksum(file_path: str) -> str:
"""Calculate MD5 checksum of a file."""

View File

@@ -297,7 +297,9 @@ class MainWindow(QMainWindow):
# Save window state before closing
self._save_window_state()
# Save annotation tab state if it exists
# Persist tab state and stop background work before exit
if hasattr(self, "training_tab"):
self.training_tab.shutdown()
if hasattr(self, "annotation_tab"):
self.annotation_tab.save_state()

View File

@@ -38,6 +38,9 @@ class AnnotationTab(QWidget):
self.current_image = None
self.current_image_path = None
self.current_image_id = None
self.current_annotations = []
# IDs of annotations currently selected on the canvas (multi-select)
self.selected_annotation_ids = []
self._setup_ui()
@@ -61,6 +64,8 @@ class AnnotationTab(QWidget):
self.annotation_canvas = AnnotationCanvasWidget()
self.annotation_canvas.zoom_changed.connect(self._on_zoom_changed)
self.annotation_canvas.annotation_drawn.connect(self._on_annotation_drawn)
# Selection of existing polylines (when tool is not in drawing mode)
self.annotation_canvas.annotation_selected.connect(self._on_annotation_selected)
canvas_layout.addWidget(self.annotation_canvas)
canvas_group.setLayout(canvas_layout)
@@ -80,24 +85,35 @@ class AnnotationTab(QWidget):
# Annotation tools section
self.annotation_tools = AnnotationToolsWidget(self.db_manager)
self.annotation_tools.pen_enabled_changed.connect(
self.annotation_canvas.set_pen_enabled
self.annotation_tools.polyline_enabled_changed.connect(
self.annotation_canvas.set_polyline_enabled
)
self.annotation_tools.pen_color_changed.connect(
self.annotation_canvas.set_pen_color
self.annotation_tools.polyline_pen_color_changed.connect(
self.annotation_canvas.set_polyline_pen_color
)
self.annotation_tools.pen_width_changed.connect(
self.annotation_canvas.set_pen_width
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.annotation_tools.process_annotations_requested.connect(
self._on_process_annotations
)
self.annotation_tools.show_annotations_requested.connect(
self._on_show_annotations
# Delete selected annotation on canvas
self.annotation_tools.delete_selected_annotation_requested.connect(
self._on_delete_selected_annotation
)
self.right_splitter.addWidget(self.annotation_tools)
@@ -180,6 +196,9 @@ class AnnotationTab(QWidget):
# 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()
@@ -217,7 +236,22 @@ class AnnotationTab(QWidget):
self._update_image_info()
def _on_annotation_drawn(self, points: list):
"""Handle when an annotation stroke is drawn."""
"""
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:
@@ -229,153 +263,260 @@ class AnnotationTab(QWidget):
)
return
logger.info(
f"Annotation drawn with {len(points)} points for class: {current_class['class_name']}"
)
# Future: Save annotation to database or export
if not points:
logger.warning("Annotation drawn with no points, ignoring")
return
def _on_class_selected(self, class_data: dict):
"""Handle when an object class is selected."""
logger.debug(f"Object class selected: {class_data['class_name']}")
# 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_annotation_selected(self, annotation_ids):
"""
Handle selection of existing annotations on the canvas.
Args:
annotation_ids: List of selected annotation IDs, or None/empty if cleared.
"""
if not annotation_ids:
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
logger.debug("Annotation selection cleared on canvas")
return
# Normalize to a unique, sorted list of integer IDs
ids = sorted({int(aid) for aid in annotation_ids if isinstance(aid, int)})
self.selected_annotation_ids = ids
self.annotation_tools.set_has_selected_annotation(bool(ids))
logger.debug(f"Annotations selected on canvas: IDs={ids}")
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'
)
# Changing the class filter invalidates any previous selection
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
# 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()
# Clear in-memory state and selection, but keep DB entries unchanged
self.current_annotations = []
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
logger.info("Cleared all annotations")
def _on_process_annotations(self):
"""Process annotations and save to database."""
# Check if we have an image loaded
if not self.current_image or not self.current_image_id:
QMessageBox.warning(
self, "No Image", "Please load an image before processing annotations."
)
return
# Get current class
current_class = self.annotation_tools.get_current_class()
if not current_class:
QMessageBox.warning(
def _on_delete_selected_annotation(self):
"""Handle deleting the currently selected annotation(s) (if any)."""
if not self.selected_annotation_ids:
QMessageBox.information(
self,
"No Class Selected",
"Please select an object class before processing annotations.",
"No Selection",
"No annotation is currently selected.",
)
return
# Compute annotation parameters asbounding boxes and polylines from annotations
parameters = self.annotation_canvas.get_annotation_parameters()
if not parameters:
QMessageBox.warning(
self,
"No Annotations",
"Please draw some annotations before processing.",
count = len(self.selected_annotation_ids)
if count == 1:
question = "Are you sure you want to delete the selected annotation?"
title = "Delete Annotation"
else:
question = (
f"Are you sure you want to delete the {count} selected annotations?"
)
return
title = "Delete Annotations"
# polyline = self.annotation_canvas.get_annotation_polyline()
for param in parameters:
bounds = param["bbox"]
polyline = param["polyline"]
try:
# Save annotation to database
annotation_id = self.db_manager.add_annotation(
image_id=self.current_image_id,
class_id=current_class["id"],
bbox=bounds,
annotator="manual",
segmentation_mask=polyline,
verified=False,
)
logger.info(
f"Saved annotation (ID: {annotation_id}) for class '{current_class['class_name']}' "
f"Bounding box: ({bounds[0]:.3f}, {bounds[1]:.3f}) to ({bounds[2]:.3f}, {bounds[3]:.3f})\n"
f"with {len(polyline)} polyline points"
)
# QMessageBox.information(
# self,
# "Success",
# f"Annotation saved successfully!\n\n"
# f"Class: {current_class['class_name']}\n"
# f"Bounding box: ({bounds[0]:.3f}, {bounds[1]:.3f}) to ({bounds[2]:.3f}, {bounds[3]:.3f})\n"
# f"Polyline points: {len(polyline)}",
# )
except Exception as e:
logger.error(f"Failed to save annotation: {e}")
QMessageBox.critical(
self, "Error", f"Failed to save annotation:\n{str(e)}"
)
# Optionally clear annotations after saving
reply = QMessageBox.question(
self,
"Clear Annotations",
"Do you want to clear the annotations to start a new one?",
title,
question,
QMessageBox.Yes | QMessageBox.No,
QMessageBox.Yes,
QMessageBox.No,
)
if reply != QMessageBox.Yes:
return
if reply == QMessageBox.Yes:
self.annotation_canvas.clear_annotations()
logger.info("Cleared annotations after saving")
failed_ids = []
try:
for ann_id in self.selected_annotation_ids:
try:
deleted = self.db_manager.delete_annotation(ann_id)
if not deleted:
failed_ids.append(ann_id)
except Exception as e:
logger.error(f"Failed to delete annotation ID {ann_id}: {e}")
failed_ids.append(ann_id)
def _on_show_annotations(self):
"""Load and display saved annotations from database."""
# Check if we have an image loaded
if not self.current_image or not self.current_image_id:
QMessageBox.warning(
self, "No Image", "Please load an image to view its annotations."
if failed_ids:
QMessageBox.warning(
self,
"Partial Failure",
"Some annotations could not be deleted:\n"
+ ", ".join(str(a) for a in failed_ids),
)
else:
logger.info(
f"Deleted {count} annotation(s): "
+ ", ".join(str(a) for a in self.selected_annotation_ids)
)
# Clear selection and reload annotations for the current image from DB
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
self._load_annotations_for_current_image()
except Exception as e:
logger.error(f"Failed to delete annotations: {e}")
QMessageBox.critical(
self,
"Error",
f"Failed to delete annotations:\n{str(e)}",
)
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()
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
return
try:
# Clear current annotations
self.annotation_canvas.clear_annotations()
# Retrieve annotations from database
annotations = self.db_manager.get_annotations_for_image(
self.current_annotations = self.db_manager.get_annotations_for_image(
self.current_image_id
)
if not annotations:
QMessageBox.information(
self, "No Annotations", "No saved annotations found for this image."
)
return
# Draw each annotation's polyline
drawn_count = 0
for ann in annotations:
if ann.get("segmentation_mask"):
polyline = ann["segmentation_mask"]
color = ann.get("class_color", "#FF0000")
# Draw the polyline
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} saved annotations from database")
QMessageBox.information(
self,
"Annotations Loaded",
f"Successfully loaded and displayed {drawn_count} annotation(s).",
)
# New annotations loaded; reset any selection
self.selected_annotation_ids = []
self.annotation_tools.set_has_selected_annotation(False)
self._redraw_annotations_for_current_filter()
except Exception as e:
logger.error(f"Failed to load annotations: {e}")
QMessageBox.critical(
self, "Error", f"Failed to load annotations:\n{str(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,
annotation_id=ann["id"],
)
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."""

File diff suppressed because it is too large Load Diff

View File

@@ -1,9 +1,10 @@
"""
Annotation canvas widget for drawing annotations on images.
Supports pen tool with color selection for manual annotation.
Currently supports polyline drawing tool with color selection for manual annotation.
"""
import numpy as np
import math
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QScrollArea
from PySide6.QtGui import (
@@ -19,25 +20,102 @@ from PySide6.QtGui import (
from PySide6.QtCore import Qt, QEvent, Signal, QPoint
from typing import Any, Dict, List, Optional, Tuple
from scipy.ndimage import binary_dilation, label, binary_fill_holes, find_objects
from skimage.measure import find_contours
from src.utils.image import Image, ImageLoadError
from src.utils.logger import get_logger
# For debugging visualization
import pylab as plt
logger = get_logger(__name__)
def perpendicular_distance(
point: Tuple[float, float],
start: Tuple[float, float],
end: Tuple[float, float],
) -> float:
"""Perpendicular distance from `point` to the line defined by `start`->`end`."""
(x, y), (x1, y1), (x2, y2) = point, start, end
dx = x2 - x1
dy = y2 - y1
if dx == 0.0 and dy == 0.0:
return math.hypot(x - x1, y - y1)
num = abs(dy * x - dx * y + x2 * y1 - y2 * x1)
den = math.hypot(dx, dy)
return num / den
def rdp(points: List[Tuple[float, float]], epsilon: float) -> List[Tuple[float, float]]:
"""
Recursive Ramer-Douglas-Peucker (RDP) polyline simplification.
Args:
points: List of (x, y) points.
epsilon: Maximum allowed perpendicular distance in pixels.
Returns:
Simplified list of (x, y) points including first and last.
"""
if len(points) <= 2:
return list(points)
start = points[0]
end = points[-1]
max_dist = -1.0
index = -1
for i in range(1, len(points) - 1):
d = perpendicular_distance(points[i], start, end)
if d > max_dist:
max_dist = d
index = i
if max_dist > epsilon:
# Recursive split
left = rdp(points[: index + 1], epsilon)
right = rdp(points[index:], epsilon)
# Concatenate but avoid duplicate at split point
return left[:-1] + right
# Keep only start and end
return [start, end]
def simplify_polyline(
points: List[Tuple[float, float]], epsilon: float
) -> List[Tuple[float, float]]:
"""
Simplify a polyline with RDP while preserving closure semantics.
If the polyline is closed (first == last), the duplicate last point is removed
before simplification and then re-added after simplification.
"""
if not points:
return []
pts = [(float(x), float(y)) for x, y in points]
closed = False
if len(pts) >= 2 and pts[0] == pts[-1]:
closed = True
pts = pts[:-1] # remove duplicate last for simplification
if len(pts) <= 2:
simplified = list(pts)
else:
simplified = rdp(pts, epsilon)
if closed and simplified:
if simplified[0] != simplified[-1]:
simplified.append(simplified[0])
return simplified
class AnnotationCanvasWidget(QWidget):
"""
Widget for displaying images and drawing annotations with pen tool.
Widget for displaying images and drawing annotations with zoom and drawing tools.
Features:
- Display images with zoom functionality
- Pen tool for drawing annotations
- Polyline tool for drawing annotations
- Configurable pen color and width
- Mouse-based drawing interface
- Zoom in/out with mouse wheel and keyboard
@@ -49,6 +127,9 @@ class AnnotationCanvasWidget(QWidget):
zoom_changed = Signal(float)
annotation_drawn = Signal(list) # List of (x, y) points in normalized coordinates
# Emitted when the user selects an existing polyline on the canvas.
# Carries the associated annotation_id (int) or None if selection is cleared
annotation_selected = Signal(object)
def __init__(self, parent=None):
"""Initialize the annotation canvas widget."""
@@ -63,13 +144,33 @@ class AnnotationCanvasWidget(QWidget):
self.zoom_step = 0.1
self.zoom_wheel_step = 0.15
# Drawing state
# Drawing / interaction state
self.is_drawing = False
self.pen_enabled = False
self.pen_color = QColor(255, 0, 0, 128) # Default red with 50% alpha
self.pen_width = 3
self.current_stroke = [] # Points in current stroke
self.all_strokes = [] # All completed strokes
self.polyline_enabled = False
self.polyline_pen_color = QColor(255, 0, 0, 128) # Default red with 50% alpha
self.polyline_pen_width = 3
self.show_bboxes: bool = True # Control visibility of bounding boxes
# Current stroke and stored polylines (in image coordinates, pixel units)
self.current_stroke: List[Tuple[float, float]] = []
self.polylines: List[List[Tuple[float, float]]] = []
self.stroke_meta: List[Dict[str, Any]] = [] # per-polyline style (color, width)
# Optional DB annotation_id for each stored polyline (None for temporary / unsaved)
self.polyline_annotation_ids: List[Optional[int]] = []
# Indices in self.polylines of the currently selected polylines (multi-select)
self.selected_polyline_indices: List[int] = []
# Stored bounding boxes in normalized coordinates (x_min, y_min, x_max, y_max)
self.bboxes: List[List[float]] = []
self.bbox_meta: List[Dict[str, Any]] = [] # per-bbox style (color, width)
# Legacy collection of strokes in normalized coordinates (kept for API compatibility)
self.all_strokes: List[dict] = []
# RDP simplification parameters (in pixels)
self.simplify_on_finish: bool = True
self.simplify_epsilon: float = 2.0
self.sample_threshold: float = 2.0 # minimum movement to sample a new point
self._setup_ui()
@@ -128,6 +229,12 @@ class AnnotationCanvasWidget(QWidget):
"""Clear all drawn annotations."""
self.all_strokes = []
self.current_stroke = []
self.polylines = []
self.stroke_meta = []
self.polyline_annotation_ids = []
self.selected_polyline_indices = []
self.bboxes = []
self.bbox_meta = []
self.is_drawing = False
if self.annotation_pixmap:
self.annotation_pixmap.fill(Qt.transparent)
@@ -218,21 +325,21 @@ class AnnotationCanvasWidget(QWidget):
"""Update display after drawing."""
self._apply_zoom()
def set_pen_enabled(self, enabled: bool):
"""Enable or disable pen tool."""
self.pen_enabled = enabled
def set_polyline_enabled(self, enabled: bool):
"""Enable or disable polyline tool."""
self.polyline_enabled = enabled
if enabled:
self.canvas_label.setCursor(Qt.CrossCursor)
else:
self.canvas_label.setCursor(Qt.ArrowCursor)
def set_pen_color(self, color: QColor):
"""Set pen color."""
self.pen_color = color
def set_polyline_pen_color(self, color: QColor):
"""Set polyline pen color."""
self.polyline_pen_color = color
def set_pen_width(self, width: int):
"""Set pen width."""
self.pen_width = max(1, width)
def set_polyline_pen_width(self, width: int):
"""Set polyline pen width."""
self.polyline_pen_width = max(1, width)
def get_zoom_percentage(self) -> int:
"""Get current zoom level as percentage."""
@@ -291,6 +398,41 @@ class AnnotationCanvasWidget(QWidget):
return (int(x), int(y))
return None
def _find_polyline_at(
self, img_x: float, img_y: float, threshold_px: float = 5.0
) -> Optional[int]:
"""
Find index of polyline whose geometry is within threshold_px of (img_x, img_y).
Returns the index in self.polylines, or None if none is close enough.
"""
best_index: Optional[int] = None
best_dist: float = float("inf")
for idx, polyline in enumerate(self.polylines):
if len(polyline) < 2:
continue
# Quick bounding-box check to skip obviously distant polylines
xs = [p[0] for p in polyline]
ys = [p[1] for p in polyline]
if img_x < min(xs) - threshold_px or img_x > max(xs) + threshold_px:
continue
if img_y < min(ys) - threshold_px or img_y > max(ys) + threshold_px:
continue
# Precise distance to all segments
for (x1, y1), (x2, y2) in zip(polyline[:-1], polyline[1:]):
d = perpendicular_distance(
(img_x, img_y), (float(x1), float(y1)), (float(x2), float(y2))
)
if d < best_dist:
best_dist = d
best_index = idx
if best_index is not None and best_dist <= threshold_px:
return best_index
return None
def _image_to_normalized_coords(self, x: int, y: int) -> Tuple[float, float]:
"""Convert image coordinates to normalized coordinates (0-1)."""
if self.original_pixmap is None:
@@ -300,26 +442,156 @@ class AnnotationCanvasWidget(QWidget):
norm_y = y / self.original_pixmap.height()
return (norm_x, norm_y)
def _add_polyline(
self,
img_points: List[Tuple[float, float]],
color: QColor,
width: int,
annotation_id: Optional[int] = None,
):
"""Store a polyline in image coordinates and redraw annotations."""
if not img_points or len(img_points) < 2:
return
# Ensure all points are tuples of floats
normalized_points = [(float(x), float(y)) for x, y in img_points]
self.polylines.append(normalized_points)
self.stroke_meta.append({"color": QColor(color), "width": int(width)})
self.polyline_annotation_ids.append(annotation_id)
self._redraw_annotations()
def _redraw_annotations(self):
"""Redraw all stored polylines and (optionally) bounding boxes onto the annotation pixmap."""
if self.annotation_pixmap is None:
return
# Clear existing overlay
self.annotation_pixmap.fill(Qt.transparent)
painter = QPainter(self.annotation_pixmap)
# Draw polylines
for idx, (polyline, meta) in enumerate(zip(self.polylines, self.stroke_meta)):
pen_color: QColor = meta.get("color", self.polyline_pen_color)
width: int = meta.get("width", self.polyline_pen_width)
if idx in self.selected_polyline_indices:
# Highlight selected polylines in a distinct color / width
highlight_color = QColor(255, 255, 0, 200) # yellow, semi-opaque
pen = QPen(
highlight_color,
width + 1,
Qt.SolidLine,
Qt.RoundCap,
Qt.RoundJoin,
)
else:
pen = QPen(
pen_color,
width,
Qt.SolidLine,
Qt.RoundCap,
Qt.RoundJoin,
)
painter.setPen(pen)
for (x1, y1), (x2, y2) in zip(polyline[:-1], polyline[1:]):
painter.drawLine(int(x1), int(y1), int(x2), int(y2))
# Draw bounding boxes (dashed) if enabled
if self.show_bboxes and self.original_pixmap is not None and self.bboxes:
img_width = float(self.original_pixmap.width())
img_height = float(self.original_pixmap.height())
for bbox, meta in zip(self.bboxes, self.bbox_meta):
if len(bbox) != 4:
continue
x_min_norm, y_min_norm, x_max_norm, y_max_norm = bbox
x_min = int(x_min_norm * img_width)
y_min = int(y_min_norm * img_height)
x_max = int(x_max_norm * img_width)
y_max = int(y_max_norm * img_height)
rect_width = x_max - x_min
rect_height = y_max - y_min
pen_color: QColor = meta.get("color", QColor(255, 0, 0, 128))
width: int = meta.get("width", self.polyline_pen_width)
pen = QPen(
pen_color,
width,
Qt.DashLine,
Qt.SquareCap,
Qt.MiterJoin,
)
painter.setPen(pen)
painter.drawRect(x_min, y_min, rect_width, rect_height)
painter.end()
self._update_display()
def mousePressEvent(self, event: QMouseEvent):
"""Handle mouse press events for drawing."""
if not self.pen_enabled or self.annotation_pixmap is None:
"""Handle mouse press events for drawing and selecting polylines."""
if self.annotation_pixmap is None:
super().mousePressEvent(event)
return
if event.button() == Qt.LeftButton:
# Get accurate position using global coordinates
label_pos = self.canvas_label.mapFromGlobal(event.globalPos())
img_coords = self._canvas_to_image_coords(label_pos)
# Map click to image coordinates
label_pos = self.canvas_label.mapFromGlobal(event.globalPos())
img_coords = self._canvas_to_image_coords(label_pos)
# Left button + drawing tool enabled -> start a new stroke
if event.button() == Qt.LeftButton and self.polyline_enabled:
if img_coords:
self.is_drawing = True
self.current_stroke = [img_coords]
self.current_stroke = [(float(img_coords[0]), float(img_coords[1]))]
return
# Left button + drawing tool disabled -> attempt selection of existing polyline
if event.button() == Qt.LeftButton and not self.polyline_enabled:
if img_coords:
idx = self._find_polyline_at(float(img_coords[0]), float(img_coords[1]))
if idx is not None:
if event.modifiers() & Qt.ShiftModifier:
# Multi-select mode: add to current selection (if not already selected)
if idx not in self.selected_polyline_indices:
self.selected_polyline_indices.append(idx)
else:
# Single-select mode: replace current selection
self.selected_polyline_indices = [idx]
# Build list of selected annotation IDs (ignore None entries)
selected_ids: List[int] = []
for sel_idx in self.selected_polyline_indices:
if 0 <= sel_idx < len(self.polyline_annotation_ids):
ann_id = self.polyline_annotation_ids[sel_idx]
if isinstance(ann_id, int):
selected_ids.append(ann_id)
if selected_ids:
self.annotation_selected.emit(selected_ids)
else:
# No valid DB-backed annotations in selection
self.annotation_selected.emit(None)
else:
# Clicked on empty space -> clear selection
self.selected_polyline_indices = []
self.annotation_selected.emit(None)
self._redraw_annotations()
return
# Fallback for other buttons / cases
super().mousePressEvent(event)
def mouseMoveEvent(self, event: QMouseEvent):
"""Handle mouse move events for drawing."""
if (
not self.is_drawing
or not self.pen_enabled
or not self.polyline_enabled
or self.annotation_pixmap is None
):
super().mouseMoveEvent(event)
@@ -330,18 +602,33 @@ class AnnotationCanvasWidget(QWidget):
img_coords = self._canvas_to_image_coords(label_pos)
if img_coords and len(self.current_stroke) > 0:
# Draw line from last point to current point
last_point = self.current_stroke[-1]
dx = img_coords[0] - last_point[0]
dy = img_coords[1] - last_point[1]
# Only sample a new point if we moved enough pixels
if math.hypot(dx, dy) < self.sample_threshold:
return
# Draw line from last point to current point for interactive feedback
painter = QPainter(self.annotation_pixmap)
pen = QPen(
self.pen_color, self.pen_width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin
self.polyline_pen_color,
self.polyline_pen_width,
Qt.SolidLine,
Qt.RoundCap,
Qt.RoundJoin,
)
painter.setPen(pen)
last_point = self.current_stroke[-1]
painter.drawLine(last_point[0], last_point[1], img_coords[0], img_coords[1])
painter.drawLine(
int(last_point[0]),
int(last_point[1]),
int(img_coords[0]),
int(img_coords[1]),
)
painter.end()
self.current_stroke.append(img_coords)
self.current_stroke.append((float(img_coords[0]), float(img_coords[1])))
self._update_display()
def mouseReleaseEvent(self, event: QMouseEvent):
@@ -352,23 +639,44 @@ class AnnotationCanvasWidget(QWidget):
self.is_drawing = False
if len(self.current_stroke) > 1:
# Convert to normalized coordinates and save stroke
normalized_stroke = [
self._image_to_normalized_coords(x, y) for x, y in self.current_stroke
]
self.all_strokes.append(
{
"points": normalized_stroke,
"color": self.pen_color.name(),
"alpha": self.pen_color.alpha(),
"width": self.pen_width,
}
)
if len(self.current_stroke) > 1 and self.original_pixmap is not None:
# Ensure the stroke is closed by connecting end -> start
raw_points = list(self.current_stroke)
if raw_points[0] != raw_points[-1]:
raw_points.append(raw_points[0])
# Emit signal with normalized coordinates
self.annotation_drawn.emit(normalized_stroke)
logger.debug(f"Completed stroke with {len(normalized_stroke)} points")
# Optional RDP simplification (in image pixel space)
if self.simplify_on_finish:
simplified = simplify_polyline(raw_points, self.simplify_epsilon)
else:
simplified = raw_points
if len(simplified) >= 2:
# Store polyline and redraw all annotations
self._add_polyline(
simplified, self.polyline_pen_color, self.polyline_pen_width
)
# Convert to normalized coordinates for metadata + signal
normalized_stroke = [
self._image_to_normalized_coords(int(x), int(y))
for (x, y) in simplified
]
self.all_strokes.append(
{
"points": normalized_stroke,
"color": self.polyline_pen_color.name(),
"alpha": self.polyline_pen_color.alpha(),
"width": self.polyline_pen_width,
}
)
# Emit signal with normalized coordinates
self.annotation_drawn.emit(normalized_stroke)
logger.debug(
f"Completed stroke with {len(simplified)} points "
f"(normalized len={len(normalized_stroke)})"
)
self.current_stroke = []
@@ -376,152 +684,61 @@ class AnnotationCanvasWidget(QWidget):
"""Get all drawn strokes with metadata."""
return self.all_strokes
# def get_annotation_bounds(self) -> Optional[Tuple[float, float, float, float]]:
# """
# Compute bounding box that encompasses all annotation strokes.
# Returns:
# Tuple of (x_min, y_min, x_max, y_max) in normalized coordinates (0-1),
# or None if no annotations exist.
# """
# if not self.all_strokes:
# return None
# # Find min/max across all strokes
# all_x = []
# all_y = []
# for stroke in self.all_strokes:
# for x, y in stroke["points"]:
# all_x.append(x)
# all_y.append(y)
# if not all_x:
# return None
# x_min = min(all_x)
# y_min = min(all_y)
# x_max = max(all_x)
# y_max = max(all_y)
# return (x_min, y_min, x_max, y_max)
# def get_annotation_polyline(self) -> List[List[float]]:
# """
# Get polyline coordinates representing all annotation strokes.
# Returns:
# List of [x, y] coordinate pairs in normalized coordinates (0-1).
# """
# polyline = []
# fig = plt.figure()
# ax1 = fig.add_subplot(411)
# ax2 = fig.add_subplot(412)
# ax3 = fig.add_subplot(413)
# ax4 = fig.add_subplot(414)
# # Get np.arrays from annotation_pixmap accoriding to the color of the stroke
# qimage = self.annotation_pixmap.toImage()
# arr = np.ndarray(
# (qimage.height(), qimage.width(), 4),
# buffer=qimage.constBits(),
# strides=[qimage.bytesPerLine(), 4, 1],
# dtype=np.uint8,
# )
# print(arr.shape, arr.dtype, arr.min(), arr.max())
# arr = np.sum(arr, axis=2)
# ax1.imshow(arr)
# arr_bin = arr > 0
# ax2.imshow(arr_bin)
# arr_bin = binary_fill_holes(arr_bin)
# ax3.imshow(arr_bin)
# labels, _number_of_features = label(
# arr_bin,
# )
# ax4.imshow(labels)
# objects = find_objects(labels)
# bounding_boxes = np.array(
# [[obj[0].start, obj[0].stop, obj[1].start, obj[1].stop] for obj in objects]
# ) / np.array([arr.shape[0], arr.shape[1]])
# print(objects)
# print(bounding_boxes)
# print(np.array([arr.shape[0], arr.shape[1]]))
# polylines = find_contours(arr_bin, 0.5)
# for pl in polylines:
# ax1.plot(pl[:, 1], pl[:, 0], "k")
# print(arr.shape, arr.dtype, arr.min(), arr.max())
# plt.show()
# return polyline
def get_annotation_parameters(self) -> Dict[str, Any]:
def get_annotation_parameters(self) -> Optional[List[Dict[str, Any]]]:
"""
Get all annotation parameters including bounding box and polyline.
Returns:
Dictionary containing:
- 'bbox': Bounding box coordinates (x_min, y_min, x_max, y_max)
- 'polyline': List of [x, y] coordinate pairs
List of dictionaries, each containing:
- 'bbox': [x_min, y_min, x_max, y_max] in normalized image coordinates
- 'polyline': List of [y_norm, x_norm] points describing the polygon
"""
# Get np.arrays from annotation_pixmap accoriding to the color of the stroke
qimage = self.annotation_pixmap.toImage()
arr = np.ndarray(
(qimage.height(), qimage.width(), 4),
buffer=qimage.constBits(),
strides=[qimage.bytesPerLine(), 4, 1],
dtype=np.uint8,
)
arr = np.sum(arr, axis=2)
arr_bin = arr > 0
arr_bin = binary_fill_holes(arr_bin)
labels, _number_of_features = label(
arr_bin,
)
if _number_of_features == 0:
if self.original_pixmap is None or not self.polylines:
return None
objects = find_objects(labels)
w, h = arr.shape
bounding_boxes = [
[obj[0].start / w, obj[1].start / h, obj[0].stop / w, obj[1].stop / h]
for obj in objects
]
img_width = float(self.original_pixmap.width())
img_height = float(self.original_pixmap.height())
polylines = find_contours(arr_bin, 0.5)
params = []
for i, pl in enumerate(polylines):
# pl is in [row, col] format from find_contours
# We need to normalize: row/height, col/width
# w = height (rows), h = width (cols) from line 510
normalized_polyline = (pl[::-1] / np.array([w, h])).tolist()
params: List[Dict[str, Any]] = []
logger.debug(f"Polyline {i}: {len(pl)} points")
logger.debug(f" w={w} (height), h={h} (width)")
logger.debug(f" First 3 normalized points: {normalized_polyline[:3]}")
for idx, polyline in enumerate(self.polylines):
if len(polyline) < 2:
continue
xs = [p[0] for p in polyline]
ys = [p[1] for p in polyline]
x_min_norm = min(xs) / img_width
x_max_norm = max(xs) / img_width
y_min_norm = min(ys) / img_height
y_max_norm = max(ys) / img_height
# Store polyline as [y_norm, x_norm] to match DB convention and
# the expectations of draw_saved_polyline().
normalized_polyline = [
[y / img_height, x / img_width] for (x, y) in polyline
]
logger.debug(
f"Polyline {idx}: {len(polyline)} points, "
f"bbox=({x_min_norm:.3f}, {y_min_norm:.3f})-({x_max_norm:.3f}, {y_max_norm:.3f})"
)
params.append(
{
"bbox": bounding_boxes[i],
"bbox": [x_min_norm, y_min_norm, x_max_norm, y_max_norm],
"polyline": normalized_polyline,
}
)
return params
return params or None
def draw_saved_polyline(
self, polyline: List[List[float]], color: str, width: int = 3
self,
polyline: List[List[float]],
color: str,
width: int = 3,
annotation_id: Optional[int] = None,
):
"""
Draw a polyline from database coordinates onto the annotation canvas.
@@ -548,36 +765,24 @@ class AnnotationCanvasWidget(QWidget):
logger.debug(f" Image size: {img_width}x{img_height}")
logger.debug(f" First 3 normalized points from DB: {polyline[:3]}")
img_coords = []
img_coords: List[Tuple[float, float]] = []
for y_norm, x_norm in polyline:
x = int(x_norm * img_width)
y = int(y_norm * img_height)
x = float(x_norm * img_width)
y = float(y_norm * img_height)
img_coords.append((x, y))
logger.debug(f" First 3 pixel coords: {img_coords[:3]}")
# Draw polyline on annotation pixmap
painter = QPainter(self.annotation_pixmap)
# Store and redraw using common pipeline
pen_color = QColor(color)
pen_color.setAlpha(128) # Add semi-transparency
pen = QPen(pen_color, width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
painter.setPen(pen)
self._add_polyline(img_coords, pen_color, width, annotation_id=annotation_id)
# Draw lines between consecutive points
for i in range(len(img_coords) - 1):
x1, y1 = img_coords[i]
x2, y2 = img_coords[i + 1]
painter.drawLine(x1, y1, x2, y2)
painter.end()
# Store in all_strokes for consistency
# Store in all_strokes for consistency (uses normalized coordinates)
self.all_strokes.append(
{"points": polyline, "color": color, "alpha": 128, "width": width}
)
# Update display
self._update_display()
logger.debug(
f"Drew saved polyline with {len(polyline)} points in color {color}"
)
@@ -587,7 +792,7 @@ class AnnotationCanvasWidget(QWidget):
Draw a bounding box from database coordinates onto the annotation canvas.
Args:
bbox: Bounding box as [y_min_norm, x_min_norm, y_max_norm, x_max_norm]
bbox: Bounding box as [x_min_norm, y_min_norm, x_max_norm, y_max_norm]
in normalized coordinates (0-1)
color: Color hex string (e.g., '#FF0000')
width: Line width in pixels
@@ -602,12 +807,11 @@ class AnnotationCanvasWidget(QWidget):
)
return
# Convert normalized coordinates to image coordinates
# bbox format: [y_min_norm, x_min_norm, y_max_norm, x_max_norm]
# Convert normalized coordinates to image coordinates (for logging/debug)
img_width = self.original_pixmap.width()
img_height = self.original_pixmap.height()
y_min_norm, x_min_norm, y_max_norm, x_max_norm = bbox
x_min_norm, y_min_norm, x_max_norm, y_max_norm = bbox
x_min = int(x_min_norm * img_width)
y_min = int(y_min_norm * img_height)
x_max = int(x_max_norm * img_width)
@@ -617,29 +821,35 @@ class AnnotationCanvasWidget(QWidget):
logger.debug(f" Image size: {img_width}x{img_height}")
logger.debug(f" Pixel coords: ({x_min}, {y_min}) to ({x_max}, {y_max})")
# Draw bounding box on annotation pixmap
painter = QPainter(self.annotation_pixmap)
# Store bounding box (normalized) and its style; actual drawing happens
# in _redraw_annotations() together with all polylines.
pen_color = QColor(color)
pen_color.setAlpha(128) # Add semi-transparency
pen = QPen(pen_color, width, Qt.SolidLine, Qt.SquareCap, Qt.MiterJoin)
painter.setPen(pen)
# Draw rectangle
rect_width = x_max - x_min
rect_height = y_max - y_min
painter.drawRect(x_min, y_min, rect_width, rect_height)
painter.end()
self.bboxes.append(
[float(x_min_norm), float(y_min_norm), float(x_max_norm), float(y_max_norm)]
)
self.bbox_meta.append({"color": pen_color, "width": int(width)})
# Store in all_strokes for consistency
self.all_strokes.append(
{"bbox": bbox, "color": color, "alpha": 128, "width": width}
)
# Update display
self._update_display()
# Redraw overlay (polylines + all bounding boxes)
self._redraw_annotations()
logger.debug(f"Drew saved bounding box in color {color}")
def set_show_bboxes(self, show: bool):
"""
Enable or disable drawing of bounding boxes.
Args:
show: If True, draw bounding boxes; if False, hide them.
"""
self.show_bboxes = bool(show)
logger.debug(f"Set show_bboxes to {self.show_bboxes}")
self._redraw_annotations()
def keyPressEvent(self, event: QKeyEvent):
"""Handle keyboard events for zooming."""
if event.key() in (Qt.Key_Plus, Qt.Key_Equal):

View File

@@ -1,6 +1,6 @@
"""
Annotation tools widget for controlling annotation parameters.
Includes pen tool, color picker, class selection, and annotation management.
Includes polyline tool, color picker, class selection, and annotation management.
"""
from PySide6.QtWidgets import (
@@ -12,6 +12,8 @@ from PySide6.QtWidgets import (
QPushButton,
QComboBox,
QSpinBox,
QDoubleSpinBox,
QCheckBox,
QColorDialog,
QInputDialog,
QMessageBox,
@@ -31,28 +33,33 @@ class AnnotationToolsWidget(QWidget):
Widget for annotation tool controls.
Features:
- Enable/disable pen tool
- Color selection for pen
- Enable/disable polyline tool
- Color selection for polyline 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)
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
"""
pen_enabled_changed = Signal(bool)
pen_color_changed = Signal(QColor)
pen_width_changed = Signal(int)
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()
process_annotations_requested = Signal()
show_annotations_requested = Signal()
# Request deletion of the currently selected annotation on the canvas
delete_selected_annotation_requested = Signal()
def __init__(self, db_manager: DatabaseManager, parent=None):
"""
@@ -64,7 +71,7 @@ class AnnotationToolsWidget(QWidget):
"""
super().__init__(parent)
self.db_manager = db_manager
self.pen_enabled = False
self.polyline_enabled = False
self.current_color = QColor(255, 0, 0, 128) # Red with 50% alpha
self.current_class = None
@@ -75,43 +82,51 @@ class AnnotationToolsWidget(QWidget):
"""Setup user interface."""
layout = QVBoxLayout()
# Pen Tool Group
pen_group = QGroupBox("Pen Tool")
pen_layout = QVBoxLayout()
# Polyline Tool Group
polyline_group = QGroupBox("Polyline Tool")
polyline_layout = QVBoxLayout()
# Enable/Disable pen
# Enable/Disable polyline tool
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)
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)
# Pen width control
# Polyline 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)
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()
pen_layout.addLayout(width_layout)
polyline_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)
pen_group.setLayout(pen_layout)
layout.addWidget(pen_group)
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")
@@ -122,7 +137,7 @@ class AnnotationToolsWidget(QWidget):
self.class_combo.currentIndexChanged.connect(self._on_class_selected)
class_layout.addWidget(self.class_combo)
# Add class button
# 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)
@@ -133,6 +148,17 @@ class AnnotationToolsWidget(QWidget):
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)
@@ -148,24 +174,22 @@ class AnnotationToolsWidget(QWidget):
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.show_btn = QPushButton("Show Saved Annotations")
self.show_btn.clicked.connect(self._on_show_annotations)
self.show_btn.setStyleSheet(
"QPushButton { background-color: #4CAF50; color: white; }"
)
actions_layout.addWidget(self.show_btn)
# 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)
# Delete currently selected annotation (enabled when a selection exists)
self.delete_selected_btn = QPushButton("Delete Selected Annotation")
self.delete_selected_btn.clicked.connect(self._on_delete_selected_annotation)
self.delete_selected_btn.setEnabled(False)
actions_layout.addWidget(self.delete_selected_btn)
actions_group.setLayout(actions_layout)
layout.addWidget(actions_group)
@@ -193,7 +217,7 @@ class AnnotationToolsWidget(QWidget):
# Clear and repopulate combo box
self.class_combo.clear()
self.class_combo.addItem("-- Select Class --", None)
self.class_combo.addItem("-- Select Class / Show All --", None)
for cls in classes:
self.class_combo.addItem(cls["class_name"], cls)
@@ -206,46 +230,115 @@ class AnnotationToolsWidget(QWidget):
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
def _on_polyline_toggle(self, checked: bool):
"""Handle polyline tool enable/disable."""
self.polyline_enabled = checked
if checked:
self.pen_toggle_btn.setText("Disable Pen")
self.pen_toggle_btn.setStyleSheet(
self.polyline_toggle_btn.setText("Stop Drawing Polyline")
self.polyline_toggle_btn.setStyleSheet(
"QPushButton { background-color: #4CAF50; }"
)
else:
self.pen_toggle_btn.setText("Enable Pen")
self.pen_toggle_btn.setStyleSheet("")
self.polyline_toggle_btn.setText("Start Drawing Polyline")
self.polyline_toggle_btn.setStyleSheet("")
self.pen_enabled_changed.emit(self.pen_enabled)
logger.debug(f"Pen tool {'enabled' if checked else 'disabled'}")
self.polyline_enabled_changed.emit(self.polyline_enabled)
logger.debug(f"Polyline 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_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 with alpha support."""
"""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(
self.current_color,
base_color,
self,
"Select Pen Color",
QColorDialog.ShowAlphaChannel, # Enable alpha channel selection
"Select Class Color",
QColorDialog.ShowAlphaChannel, # Allow alpha in UI, but store RGB in DB
)
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()}"
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."""
"""Handle object class selection (including '-- Select Class --')."""
class_data = self.class_combo.currentData()
if class_data:
@@ -260,20 +353,23 @@ class AnnotationToolsWidget(QWidget):
self.class_info_label.setText(info_text)
# Update pen color to match class color with semi-transparency
# 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.pen_color_changed.emit(class_color)
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."""
@@ -351,36 +447,32 @@ class AnnotationToolsWidget(QWidget):
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
def _on_delete_selected_annotation(self):
"""Handle delete selected annotation button."""
self.delete_selected_annotation_requested.emit()
logger.debug("Delete selected annotation requested")
self.process_annotations_requested.emit()
logger.debug("Process annotations requested")
def set_has_selected_annotation(self, has_selection: bool):
"""
Enable/disable actions that require a selected annotation.
def _on_show_annotations(self):
"""Handle show annotations button."""
self.show_annotations_requested.emit()
logger.debug("Show annotations requested")
Args:
has_selection: True if an annotation is currently selected on the canvas.
"""
self.delete_selected_btn.setEnabled(bool(has_selection))
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."""
def get_polyline_pen_color(self) -> QColor:
"""Get current polyline pen color."""
return self.current_color
def get_pen_width(self) -> int:
"""Get current pen width."""
return self.pen_width_spin.value()
def get_polyline_pen_width(self) -> int:
"""Get current polyline pen width."""
return self.polyline_pen_width_spin.value()
def is_pen_enabled(self) -> bool:
"""Check if pen tool is enabled."""
return self.pen_enabled
def is_polyline_enabled(self) -> bool:
"""Check if polyline tool is enabled."""
return self.polyline_enabled

View File

@@ -55,6 +55,7 @@ class YOLOWrapper:
save_dir: str = "data/models",
name: str = "custom_model",
resume: bool = False,
callbacks: Optional[Dict[str, Callable]] = None,
**kwargs,
) -> Dict[str, Any]:
"""
@@ -69,6 +70,7 @@ class YOLOWrapper:
save_dir: Directory to save trained model
name: Name for the training run
resume: Resume training from last checkpoint
callbacks: Optional Ultralytics callback dictionary
**kwargs: Additional training arguments
Returns: