# Copyright 2025-2026 Thousand Brains Project
#
# Copyright may exist in Contributors' modifications
# and/or contributions to the work.
#
# Use of this source code is governed by the MIT
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.
from __future__ import annotations
import logging
from dataclasses import dataclass
import cv2
import numpy as np
from tbp.monty.frameworks.utils.spatial_arithmetics import (
TangentFrame,
normalize,
project_onto_tangent_plane,
)
from tbp.monty.geometry import SurfacePoseVectors
from tbp.monty.math import DEFAULT_TOLERANCE
from tbp.monty.observations import SensorObservation
logger = logging.getLogger(__name__)
DEFAULT_SURFACE_POSE_VECTORS_2D = SurfacePoseVectors(
np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
)
def _gradient_to_tangent_angle(gradient_angle: float) -> float:
"""Convert gradient direction to edge tangent direction and wrap to (0, pi].
The edge tangent is perpendicular to the gradient direction.
Args:
gradient_angle: Gradient direction in radians (any range)
Returns:
Edge tangent angle in (0, pi] radians.
"""
tangent_angle = gradient_angle + np.pi / 2
return (tangent_angle + 2 * np.pi) % (2 * np.pi)
def _angle_to_pose_2d(
angle: float,
cam_to_world: np.ndarray,
surface_normal: np.ndarray,
tangent_frame: TangentFrame,
) -> SurfacePoseVectors:
"""Build 2D pose vectors from an edge angle.
The image-space edge direction is projected onto the local tangent plane and
expressed in tangent-frame coordinates.
Args:
angle: Edge angle in radians in image coordinates (y-down), measured
from the image x-axis toward the image y-axis (i.e., toward the
bottom of the image).
cam_to_world: 4x4 camera-to-world transformation matrix.
surface_normal: Surface normal defining the local tangent plane.
tangent_frame: Orthonormal local basis used for the 2D model.
Returns:
3x3 array whose rows are [normal, edge_tangent, edge_perp].
Normal is always [0, 0, 1]; tangent and perp lie in local 2D coordinates.
"""
R = cam_to_world[:3, :3] # noqa: N806
image_x_world = R @ np.array([1.0, 0.0, 0.0])
image_y_world = R @ np.array([0.0, -1.0, 0.0])
edge_world = np.cos(angle) * image_x_world + np.sin(angle) * image_y_world
edge_tangent_world = project_onto_tangent_plane(edge_world, surface_normal)
if np.linalg.norm(edge_tangent_world) < DEFAULT_TOLERANCE:
edge_tangent_world = tangent_frame.basis_u
else:
edge_tangent_world = normalize(edge_tangent_world)
edge_tangent_u = np.dot(edge_tangent_world, tangent_frame.basis_u)
edge_tangent_v = np.dot(edge_tangent_world, tangent_frame.basis_v)
edge_tangent_2d = normalize(np.array([edge_tangent_u, edge_tangent_v, 0.0]))
return SurfacePoseVectors(
np.array(
[
[0.0, 0.0, 1.0],
edge_tangent_2d,
[-edge_tangent_2d[1], edge_tangent_2d[0], 0.0],
]
)
)
[docs]
@dataclass
class StructureTensor:
"""Structure tensor at a single point.
A 2x2 symmetric matrix [[Jxx, Jxy], [Jxy, Jyy]].
"""
xx: float
yy: float
xy: float
@property
def eigenvalues(self) -> tuple[float, float]:
"""Returns (lambda_min, lambda_max) of the 2x2 structure tensor."""
matrix = np.array([[self.xx, self.xy], [self.xy, self.yy]])
lambda_min, lambda_max = np.linalg.eigh(matrix)[0]
return lambda_min, lambda_max
@property
def gradient_theta(self) -> float:
"""Gradient direction in radians (normal to the dominant edge)."""
return 0.5 * np.arctan2(2.0 * self.xy, self.xx - self.yy)
@property
def edge_strength(self) -> float:
"""Magnitude of the dominant eigenvalue."""
_, lambda_max = self.eigenvalues
return np.sqrt(max(lambda_max, 0.0))
@property
def coherence(self) -> float:
"""Edge quality in [0, 1]: 1 means perfectly oriented, 0 means isotropic."""
lambda_min, lambda_max = self.eigenvalues
lambda_diff = lambda_max - lambda_min
if lambda_diff < DEFAULT_TOLERANCE:
return 0.0
return np.clip(lambda_diff / (lambda_max + lambda_min), 0.0, 1.0)
@property
def edge_angle(self) -> float:
"""Edge angle in [0, pi] radians."""
return _gradient_to_tangent_angle(self.gradient_theta)
[docs]
@dataclass
class EdgeFeatures:
"""Edge features extracted from a single image patch."""
angle: float | None
strength: float
coherence: float
is_geometric_edge: bool
has_edge: bool
[docs]
class EdgeDetector:
[docs]
def __init__(
self,
gaussian_sigma: float = 1.0,
kernel_size: int = 7,
strength_threshold: float = 0.1,
coherence_threshold: float = 0.5,
radius: float = 14.0,
sigma_r: float = 7.0,
depth_edge_threshold: float = 0.01,
max_center_offset: int | None = None,
):
"""Initialize EdgeDetector.
Args:
gaussian_sigma: float = 1.0
kernel_size: int = 7
strength_threshold: float = 0.1
coherence_threshold: float = 0.5
radius: float = 14.0
sigma_r: float = 7.0
depth_edge_threshold: float = 0.01
max_center_offset: int | None = None
Raises:
ValueError: Invalid argument for gaussian_sigma.
"""
self._gaussian_sigma = gaussian_sigma
self._kernel_size = kernel_size
self._strength_threshold = strength_threshold
self._coherence_threshold = coherence_threshold
self._radius = radius
self._sigma_r = sigma_r
self._depth_edge_threshold = depth_edge_threshold
self._max_center_offset = max_center_offset
if self._kernel_size % 2 != 1:
raise ValueError("Kernel size must be odd.")
def __call__(
self,
observation: SensorObservation,
) -> EdgeFeatures:
"""Compute edge features using center-weighted, global-aware structure tensor.
This function aggregates structure tensor components over a center-biased
neighborhood, giving higher weight to pixels closer to the center and pixels
with stronger gradients.
Reference:
Nazar Khan, "Corner Detection" lecture notes, Section on Structure
Tensor. http://faculty.pucit.edu.pk/nazarkhan/teaching/Spring2021/CS565/Lectures/lecture6_corner_detection.pdf
Args:
observation: Sensor observation.
Returns:
EdgeFeatures.
"""
patch = observation["rgba"][:, :, :3]
depth = observation["depth"]
grayscale = cv2.cvtColor(patch, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
Ix, Iy = self._compute_sobel_gradients(grayscale) # noqa: N806
tensor_per_pixel = self._compute_per_pixel_structure_tensors(Ix, Iy)
weights, total_weight = self._compute_center_weights(grayscale.shape, Ix, Iy)
if total_weight < DEFAULT_TOLERANCE:
return EdgeFeatures(
angle=None,
strength=0.0,
coherence=0.0,
is_geometric_edge=False,
has_edge=False,
)
aggregated = self._aggregate_tensor(tensor_per_pixel, weights, total_weight)
edge_normal_angle = aggregated.edge_angle + np.pi / 2
edge_center_offset = self._compute_center_offset(
weights,
total_weight,
edge_normal_angle,
)
if not self._passes_center_check(edge_center_offset):
return EdgeFeatures(
angle=None,
strength=0.0,
coherence=0.0,
is_geometric_edge=False,
has_edge=False,
)
has_edge = (
aggregated.edge_strength > self._strength_threshold
and aggregated.coherence > self._coherence_threshold
)
return EdgeFeatures(
angle=aggregated.edge_angle,
strength=aggregated.edge_strength,
coherence=aggregated.coherence,
is_geometric_edge=self._is_geometric_edge(
depth,
aggregated.edge_angle,
edge_center_offset,
),
has_edge=has_edge,
)
@staticmethod
def _compute_sobel_gradients(
grayscale: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""Compute horizontal and vertical Sobel gradients.
Args:
grayscale: Grayscale image patch as float32 in [0, 1].
Returns:
Tuple of (Ix, Iy): horizontal and vertical gradient arrays.
"""
Ix = cv2.Sobel(grayscale, cv2.CV_32F, 1, 0, ksize=3) # noqa: N806
Iy = cv2.Sobel(grayscale, cv2.CV_32F, 0, 1, ksize=3) # noqa: N806
return Ix, Iy
def _compute_per_pixel_structure_tensors(
self,
Ix: np.ndarray, # noqa: N803
Iy: np.ndarray, # noqa: N803
) -> np.ndarray:
"""Build the smoothed per-pixel structure tensor field from Sobel gradients.
Args:
Ix: Horizontal Sobel gradients.
Iy: Vertical Sobel gradients.
Returns:
Array of shape (h, w, 2, 2) where each entry is the Gaussian-smoothed
2x2 structure tensor [[Jxx, Jxy], [Jxy, Jyy]] at that pixel.
"""
Jxx = Ix * Ix # noqa: N806 # (h, w)
Jyy = Iy * Iy # noqa: N806 # (h, w)
Jxy = Ix * Iy # noqa: N806 # (h, w)
ksize = self._kernel_size
sigma = self._gaussian_sigma
Jxx = cv2.GaussianBlur(Jxx, (ksize, ksize), sigma) # noqa: N806 # (h, w)
Jyy = cv2.GaussianBlur(Jyy, (ksize, ksize), sigma) # noqa: N806 # (h, w)
Jxy = cv2.GaussianBlur(Jxy, (ksize, ksize), sigma) # noqa: N806 # (h, w)
h, w = Jxx.shape
tensor_per_pixel = np.empty((h, w, 2, 2), dtype=np.float32)
tensor_per_pixel[..., 0, 0] = Jxx
tensor_per_pixel[..., 1, 1] = Jyy
tensor_per_pixel[..., 0, 1] = Jxy
tensor_per_pixel[..., 1, 0] = Jxy
return tensor_per_pixel
def _compute_center_weights(
self,
shape: tuple[int, int],
Ix: np.ndarray, # noqa: N803
Iy: np.ndarray, # noqa: N803
) -> tuple[np.ndarray, np.floating]:
"""Build radial + gradient-strength weight map centered on the patch.
Weights combine a Gaussian radial falloff (suppressing far-from-center pixels)
with local gradient magnitude (so strong off-center edges still contribute).
Args:
shape: (height, width) of the patch.
Ix: Horizontal Sobel gradients.
Iy: Vertical Sobel gradients.
Returns:
Tuple of (weights, total_weight).
"""
h, w = shape
r0, c0 = h // 2, w // 2
rows, cols = np.meshgrid(np.arange(h), np.arange(w), indexing="ij")
d_squared = (rows - r0) ** 2 + (cols - c0) ** 2
d = np.sqrt(d_squared)
w_r = np.exp(-d_squared / (2.0 * self._sigma_r**2))
w_r[d > self._radius] = 0.0
weights = w_r * (Ix**2 + Iy**2)
total_weight = np.sum(weights)
return weights, total_weight
@staticmethod
def _aggregate_tensor(
tensor_per_pixel: np.ndarray,
weights: np.ndarray,
total_weight: float,
) -> StructureTensor:
"""Reduce a per-pixel structure tensor field to a single representative tensor.
Args:
tensor_per_pixel: Per-pixel structure tensors, shape (h, w, 2, 2).
weights: Per-pixel weights, shape (h, w).
total_weight: Sum of weights (must be > 0).
Returns:
StructureTensor representing the weighted aggregate over the patch.
"""
w = weights[..., np.newaxis, np.newaxis]
aggregated = np.sum(w * tensor_per_pixel, axis=(0, 1)) / total_weight
return StructureTensor(
xx=float(aggregated[0, 0]),
yy=float(aggregated[1, 1]),
xy=float(aggregated[0, 1]),
)
@staticmethod
def _compute_center_offset(
weights: np.ndarray,
total_weight: np.floating,
normal_angle: float,
) -> np.floating:
"""Estimate signed edge offset from the patch center along a normal.
Args:
weights: Per-pixel weights.
total_weight: Sum of weights.
normal_angle: Direction perpendicular to the edge in radians.
Returns:
Weighted signed offset from the patch center in pixels.
"""
h, w = weights.shape
r0, c0 = h // 2, w // 2
rows, cols = np.meshgrid(np.arange(h), np.arange(w), indexing="ij")
nx = np.cos(normal_angle)
ny = np.sin(normal_angle)
dist_normal = nx * (cols - c0) + ny * (rows - r0)
return np.sum(weights * dist_normal) / total_weight
def _passes_center_check(self, center_offset: float) -> bool:
"""Return True if the detected edge passes close enough to the patch center.
Args:
center_offset: Signed edge offset from the patch center in pixels.
Returns:
True if edge passes the center check (or check is disabled).
"""
if self._max_center_offset is None:
return True
return abs(center_offset) <= self._max_center_offset
def _is_geometric_edge(
self,
depth_patch: np.ndarray,
edge_theta: float,
edge_center_offset: float,
) -> bool:
"""Check if detected edge is a geometric edge (depth discontinuity).
Geometric edges occur at object boundaries or surface creases where depth
changes abruptly. Texture edges will be detected wherever there is an abrupt
discontinuity in image intensity. We will use detected geometric edges to
identify candidate texture edges that do not correspond to a 2D surface
(such as where the red handle of a mug is seen against the black background
of a simulator's void). This function computes the depth gradient perpendicular
to the detected edge direction near the estimated RGB edge line and
checks if robust depth-edge evidence exceeds a threshold.
Args:
depth_patch: Depth image patch (same size as RGB patch used for edge
detection). Values should be in consistent units (e.g., meters).
edge_theta: Edge tangent angle in radians from RGB edge detection.
edge_center_offset: Signed RGB edge offset from the patch center along
the edge normal.
Returns:
True if edge is geometric, False if texture edge.
"""
depth_dx = cv2.Sobel(depth_patch, cv2.CV_32F, 1, 0, ksize=3)
depth_dy = cv2.Sobel(depth_patch, cv2.CV_32F, 0, 1, ksize=3)
edge_normal_angle = edge_theta + np.pi / 2
nx = np.cos(edge_normal_angle)
ny = np.sin(edge_normal_angle)
h, w = depth_patch.shape
cy, cx = h // 2, w // 2
rows, cols = np.meshgrid(np.arange(h), np.arange(w), indexing="ij")
dist_normal = nx * (cols - cx) + ny * (rows - cy)
depth_edge_band = np.abs(dist_normal - edge_center_offset) <= 1.5
if not np.any(depth_edge_band):
return False
depth_gradient_perp = np.abs(nx * depth_dx + ny * depth_dy)
depth_edge_evidence = np.percentile(
depth_gradient_perp[depth_edge_band],
95.0,
)
return depth_edge_evidence > self._depth_edge_threshold
