From 47b7e896988925a37eba1d8480a6db9bb01b3f19 Mon Sep 17 00:00:00 2001 From: x01dc Date: Sun, 6 Sep 2026 13:27:56 +0200 Subject: [PATCH] feat(vision): add Layer-0 image-processing toolkit for the vision plan First implementation slice of AI_docs/MACHINE_VISION_AUTOCONFIRM_PLAN.md (Phase 0 groundwork): pure numpy/opencv functions with no device/BEC coupling -- grayscale conversion, CLAHE, Canny edge detection, ECC-based registration (never raises on non-convergence, reports converged=False instead), contour extraction/scoring, multi-reference best-match, and a diagnostic contour overlay. This is the shared building block Layer 1 (VisionInterface) and Layer 2 (SampleVisionMonitor) will be built on. 16 new tests, full suite (322 tests) still green. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_018Mupip8x7Hean6E3FFvpbj --- csaxs_bec/devices/vision/__init__.py | 5 + csaxs_bec/devices/vision/vision_toolkit.py | 174 +++++++++++++++++++++ tests/tests_devices/test_vision_toolkit.py | 138 ++++++++++++++++ 3 files changed, 317 insertions(+) create mode 100644 csaxs_bec/devices/vision/__init__.py create mode 100644 csaxs_bec/devices/vision/vision_toolkit.py create mode 100644 tests/tests_devices/test_vision_toolkit.py diff --git a/csaxs_bec/devices/vision/__init__.py b/csaxs_bec/devices/vision/__init__.py new file mode 100644 index 00000000..c1f6aadd --- /dev/null +++ b/csaxs_bec/devices/vision/__init__.py @@ -0,0 +1,5 @@ +"""Generic, camera-agnostic image-processing/vision components. + +See ``csaxs_bec/bec_ipython_client/plugins/omny/AI_docs/MACHINE_VISION_AUTOCONFIRM_PLAN.md`` +for the design this package implements (Layer 0: :mod:`vision_toolkit`). +""" diff --git a/csaxs_bec/devices/vision/vision_toolkit.py b/csaxs_bec/devices/vision/vision_toolkit.py new file mode 100644 index 00000000..e87ba025 --- /dev/null +++ b/csaxs_bec/devices/vision/vision_toolkit.py @@ -0,0 +1,174 @@ +"""Generic image-processing building blocks for camera-based visual checks. + +Pure ``numpy``/``opencv`` functions -- array in, array (or a small result object) out. No BEC or +ophyd imports, no camera-device coupling: this is Layer 0 of +``AI_docs/MACHINE_VISION_AUTOCONFIRM_PLAN.md``, meant to be shared by any device or client-side code +that needs to compare a live frame against reference images (OMNY sample-transfer confirms today, +potentially an x-ray-eye-alignment processed view later). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field + +import cv2 +import numpy as np + + +@dataclass(frozen=True) +class RegistrationResult: + """Result of aligning a frame to a reference via :func:`register_to_reference`.""" + + aligned_image: np.ndarray + warp_matrix: np.ndarray + converged: bool + + +@dataclass(frozen=True) +class ReferenceMatch: + """Score of ``current_image`` against a single reference, from :func:`best_reference_match`.""" + + reference_index: int + contour_score: float + registration: RegistrationResult + + +@dataclass(frozen=True) +class MatchResult: + """Best-match result of :func:`best_reference_match` across a whole reference set.""" + + best: ReferenceMatch | None + all_matches: list[ReferenceMatch] = field(default_factory=list) + + @property + def best_score(self) -> float: + """The lowest (best) contour score found, or +inf if there was nothing to compare against.""" + return self.best.contour_score if self.best is not None else math.inf + + +def to_grayscale(image: np.ndarray) -> np.ndarray: + """Convert an image to single-channel grayscale. + + Frames from ``get_last_image()`` (``IDSCamera``/``AlliedVisionAravisCamera``) are already + display-oriented RGB, not OpenCV's native BGR (see ``Camera.get_image_data()``'s explicit + BGR->RGB channel reshuffle) -- so this uses ``COLOR_RGB2GRAY``, not the more commonly seen + ``COLOR_BGR2GRAY``. + """ + if image.ndim == 2: + return image + if image.ndim == 3 and image.shape[2] in (3, 4): + return cv2.cvtColor(image[:, :, :3], cv2.COLOR_RGB2GRAY) + raise ValueError(f"Unsupported image shape for grayscale conversion: {image.shape}") + + +def apply_clahe( + image: np.ndarray, clip_limit: float = 2.0, tile_grid_size: tuple[int, int] = (8, 8) +) -> np.ndarray: + """Contrast-normalize a grayscale image so exposure differences don't dominate a comparison.""" + gray = to_grayscale(image) + clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size) + return clahe.apply(gray) + + +def detect_edges( + image: np.ndarray, low_threshold: int = 50, high_threshold: int = 150 +) -> np.ndarray: + """Canny edge map of a grayscale (or CLAHE-normalized) image.""" + gray = to_grayscale(image) + return cv2.Canny(gray, low_threshold, high_threshold) + + +def register_to_reference( + reference: np.ndarray, + current: np.ndarray, + warp_mode: int = cv2.MOTION_EUCLIDEAN, + max_iterations: int = 50, + epsilon: float = 1e-4, +) -> RegistrationResult: + """Align ``current`` to ``reference`` (translation+rotation) to compensate for small camera drift. + + ``cv2.findTransformECC`` raises ``cv2.error`` outright if it fails to converge (e.g. a + low-texture or badly mismatched pair) -- caught here rather than propagated, since a failed + registration should fall back to comparing the unregistered frame, not abort the caller's + confirm gate. + """ + reference_gray = to_grayscale(reference).astype(np.float32) + current_gray = to_grayscale(current).astype(np.float32) + warp_matrix = np.eye(2, 3, dtype=np.float32) + criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, max_iterations, epsilon) + try: + _, warp_matrix = cv2.findTransformECC( + reference_gray, current_gray, warp_matrix, warp_mode, criteria + ) + except cv2.error: + return RegistrationResult(aligned_image=current, warp_matrix=warp_matrix, converged=False) + + height, width = reference_gray.shape + aligned = cv2.warpAffine( + current, warp_matrix, (width, height), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP + ) + return RegistrationResult(aligned_image=aligned, warp_matrix=warp_matrix, converged=True) + + +def largest_contour(binary_image: np.ndarray) -> np.ndarray | None: + """The largest-area contour in a binary/edge image (e.g. from :func:`detect_edges`), if any.""" + contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + return None + return max(contours, key=cv2.contourArea) + + +def contour_similarity_score(contour_a: np.ndarray | None, contour_b: np.ndarray | None) -> float: + """Shape-distance between two contours (lower = more similar); +inf if either is missing.""" + if contour_a is None or contour_b is None: + return math.inf + return cv2.matchShapes(contour_a, contour_b, cv2.CONTOURS_MATCH_I1, 0.0) + + +def best_reference_match( + current_image: np.ndarray, + reference_images: list[np.ndarray], + edge_low_threshold: int = 50, + edge_high_threshold: int = 150, +) -> MatchResult: + """Register+score ``current_image`` against every image in ``reference_images``. + + Scores by contour shape (edge-based, so largely exposure-invariant) after registering each + reference to the current frame. Returns every per-reference match plus the best (lowest-score) + one -- an empty ``reference_images`` list yields ``MatchResult(best=None)``, not an error, since + "no references yet" is an expected state (see the plan's "reference corpus curation" open + question), not a caller mistake. + """ + current_edges = detect_edges(current_image, edge_low_threshold, edge_high_threshold) + current_contour = largest_contour(current_edges) + + matches: list[ReferenceMatch] = [] + for index, reference_image in enumerate(reference_images): + # Warp the *reference* into the current frame (not the other way round) so + # current_contour, computed once above from the unwarped live frame, is directly + # comparable to reference_contour below without re-deriving it per reference. + registration = register_to_reference(current_image, reference_image) + reference_edges = detect_edges( + registration.aligned_image, edge_low_threshold, edge_high_threshold + ) + reference_contour = largest_contour(reference_edges) + score = contour_similarity_score(current_contour, reference_contour) + matches.append( + ReferenceMatch(reference_index=index, contour_score=score, registration=registration) + ) + + best = min(matches, key=lambda match: match.contour_score) if matches else None + return MatchResult(best=best, all_matches=matches) + + +def draw_contour_overlay( + base_image: np.ndarray, contour: np.ndarray | None, color: tuple[int, int, int] = (255, 0, 0) +) -> np.ndarray: + """A display-ready copy of ``base_image`` with ``contour`` drawn on top, for diagnostic viewing.""" + overlay = base_image.copy() + if overlay.ndim == 2: + overlay = cv2.cvtColor(overlay, cv2.COLOR_GRAY2RGB) + if contour is not None: + cv2.drawContours(overlay, [contour], -1, color, thickness=2) + return overlay diff --git a/tests/tests_devices/test_vision_toolkit.py b/tests/tests_devices/test_vision_toolkit.py new file mode 100644 index 00000000..564e8f19 --- /dev/null +++ b/tests/tests_devices/test_vision_toolkit.py @@ -0,0 +1,138 @@ +"""Unit tests for the Layer-0 vision toolkit (csaxs_bec/devices/vision/vision_toolkit.py).""" + +import cv2 +import numpy as np +import pytest + +from csaxs_bec.devices.vision import vision_toolkit as vt + + +def _square_image( + size: int = 100, square_side: int = 30, offset: tuple[int, int] = (0, 0) +) -> np.ndarray: + """A synthetic RGB frame: a white filled square on a black background.""" + image = np.zeros((size, size, 3), dtype=np.uint8) + top_left = (size // 2 - square_side // 2 + offset[0], size // 2 - square_side // 2 + offset[1]) + bottom_right = (top_left[0] + square_side, top_left[1] + square_side) + cv2.rectangle(image, top_left, bottom_right, color=(255, 255, 255), thickness=-1) + return image + + +def test_to_grayscale_passes_through_2d(): + gray = np.zeros((10, 10), dtype=np.uint8) + assert vt.to_grayscale(gray) is gray + + +def test_to_grayscale_converts_rgb(): + image = _square_image() + gray = vt.to_grayscale(image) + assert gray.shape == (100, 100) + assert gray.dtype == np.uint8 + + +def test_to_grayscale_rejects_unsupported_shape(): + with pytest.raises(ValueError): + vt.to_grayscale(np.zeros((10,), dtype=np.uint8)) + + +def test_apply_clahe_returns_grayscale_shape(): + image = _square_image() + result = vt.apply_clahe(image) + assert result.shape == (100, 100) + assert result.dtype == np.uint8 + + +def test_detect_edges_finds_the_square_boundary(): + image = _square_image() + edges = vt.detect_edges(image) + assert edges.shape == (100, 100) + # A filled square has a boundary -- some edge pixels must be set. + assert edges.sum() > 0 + + +def test_detect_edges_blank_image_has_no_edges(): + image = np.zeros((100, 100, 3), dtype=np.uint8) + edges = vt.detect_edges(image) + assert edges.sum() == 0 + + +def test_register_to_reference_identity_for_identical_images(): + image = _square_image() + result = vt.register_to_reference(image, image) + assert result.converged + # Aligning an image to itself should be close to a no-op. + assert np.allclose(result.warp_matrix, np.eye(2, 3, dtype=np.float32), atol=1e-3) + + +def test_register_to_reference_reports_non_convergence_without_raising(): + # Two flat, textureless images give ECC a zero-gradient problem it cannot solve -- + # this must be reported via converged=False, never raise out to the caller. + flat_a = np.full((50, 50, 3), 128, dtype=np.uint8) + flat_b = np.full((50, 50, 3), 128, dtype=np.uint8) + result = vt.register_to_reference(flat_a, flat_b) + assert result.converged is False + assert result.aligned_image is flat_b + + +def test_largest_contour_finds_the_square(): + image = _square_image() + edges = vt.detect_edges(image) + contour = vt.largest_contour(edges) + assert contour is not None + assert cv2.contourArea(contour) > 0 + + +def test_largest_contour_none_for_blank_image(): + blank = np.zeros((50, 50), dtype=np.uint8) + assert vt.largest_contour(blank) is None + + +def test_contour_similarity_score_zero_for_identical_contour(): + image = _square_image() + contour = vt.largest_contour(vt.detect_edges(image)) + assert vt.contour_similarity_score(contour, contour) == pytest.approx(0.0, abs=1e-6) + + +def test_contour_similarity_score_infinite_if_missing(): + image = _square_image() + contour = vt.largest_contour(vt.detect_edges(image)) + assert vt.contour_similarity_score(contour, None) == float("inf") + assert vt.contour_similarity_score(None, None) == float("inf") + + +def test_best_reference_match_picks_the_matching_reference(): + current = _square_image(offset=(0, 0)) + matching_reference = _square_image(offset=(0, 0)) + mismatched_reference = np.zeros((100, 100, 3), dtype=np.uint8) # no square at all + + result = vt.best_reference_match(current, [mismatched_reference, matching_reference]) + + assert result.best is not None + assert result.best.reference_index == 1 + assert result.best_score < vt.contour_similarity_score( + vt.largest_contour(vt.detect_edges(current)), None + ) + + +def test_best_reference_match_empty_reference_list(): + current = _square_image() + result = vt.best_reference_match(current, []) + assert result.best is None + assert result.best_score == float("inf") + assert result.all_matches == [] + + +def test_draw_contour_overlay_draws_on_a_copy(): + image = _square_image() + contour = vt.largest_contour(vt.detect_edges(image)) + overlay = vt.draw_contour_overlay(image, contour) + assert overlay.shape == image.shape + assert overlay is not image + assert not np.array_equal(overlay, image) + + +def test_draw_contour_overlay_without_contour_returns_unmodified_copy(): + image = _square_image() + overlay = vt.draw_contour_overlay(image, None) + assert np.array_equal(overlay, image) + assert overlay is not image