depth_fm.data.adapter

adapter

DepthFM dataset adapter for MarsHiRISEDTM.

Map-style wrapper around MarsHiRISEDTM + HiRISEGeoSampler that produces DepthFM-ready training pairs::

image:  (3, H, W) float32 in [-1, 1]  — orthoimage (1ch → 3ch replicated)
dtm:    (3, H, W) float32 in [-1, 1]  — normalised elevation (1ch → 3ch)

Stereo augmentation: each __getitem__ randomly selects the left or right orthoimage as the conditioning input. Both pair with the same DTM, doubling the effective training data.

The heavy image-processing primitives (void filling, seam detection, sun vector estimation, TIN-artifact detection, topographic residual) used to live inline in this file. They are now in depth_fm.data.image_processing.* and re-exported below for the existing callers (build_litdata.py, train_lightning.py, debug_viz.py, depthfm_pipeline_diagram.py).

Usage::

from dataset.core.dtm import MarsHiRISEDTM
from dataset.sampling.sampler import HiRISEGeoSampler
from depth_fm.data.adapter import DepthFMHiRISEAdapterCached

base = MarsHiRISEDTM(root="/scratch/mars_hirise_dtm", include_ortho=True,
                     ortho_type="RED", download=True)
sampler = HiRISEGeoSampler(base, size=0.009, length=10000)
adapter = DepthFMHiRISEAdapterCached(
    base_dataset=base, sampler=sampler, resolution=512,
    dtm_normalization="relative",
    stats_path="dataset_stats/dtm/dataset_stats.json",
)

SeamResult dataclass

SeamResult(seam_score: float, ortho_score: float, dtm_score: float, cohens_d: float, best_angle_rad: float, best_y: int, best_x: int, line_length: int, num_angles: int, span: float, sparsity: float, composite_score: float, is_seam: bool, seam_heatmap: Optional[ndarray] = None, cohens_d_heatmap: Optional[ndarray] = None, per_angle_max: Optional[ndarray] = None, diag_hot_mask: Optional[ndarray] = None, diag_closed_components: Optional[ndarray] = None, diag_hough_lines: Optional[ndarray] = None, diag_isolation_profile: Optional[tuple] = None)

line_endpoints

line_endpoints(clip_hw=None)

Return (y1, x1, y2, x2) pixel endpoints of the best-scoring line.

Source code in src/depth_fm/data/image_processing/seam_detection.py

def line_endpoints(self, clip_hw=None):
    """Return (y1, x1, y2, x2) pixel endpoints of the best-scoring line."""
    half = self.line_length // 2
    dx, dy = math.cos(self.best_angle_rad), math.sin(self.best_angle_rad)
    y1 = self.best_y - half * dy
    x1 = self.best_x - half * dx
    y2 = self.best_y + half * dy
    x2 = self.best_x + half * dx
    if clip_hw is not None:
        H, W = clip_hw
        y1 = float(np.clip(y1, 0, H - 1))
        y2 = float(np.clip(y2, 0, H - 1))
        x1 = float(np.clip(x1, 0, W - 1))
        x2 = float(np.clip(x2, 0, W - 1))
    return y1, x1, y2, x2

DepthFMHiRISEAdapterCached

DepthFMHiRISEAdapterCached(base_dataset: MarsHiRISEDTM, sampler: HiRISEGeoSampler, resolution: int = 512, dtm_normalization: Literal['relative', 'log', 'linear'] = 'relative', random_flip: bool = True, random_jitter: bool = False, brightness_jitter: float = 0.1, stats_path: str | None = None, clip: bool = False, use_manifest: bool = True, manifest_workers: int = 16, manifest_dir: str | None = None, erode_radius: int = 2, multiprocessing_context='fork')

Bases: Dataset

Adapter that converts MarsHiRISEDTM samples into DepthFM training pairs.

Each __getitem__ draws a geo-slice from the sampler, loads elevation + orthoimage(s) via the base dataset, runs the manifest-cached preprocessing (sun vector, valid mask, void fill), and returns a normalised dict ready for the DepthFM training loop.

Parameters:

Name	Type	Description	Default
base_dataset	MarsHiRISEDTM	An initialised MarsHiRISEDTM instance.	required
sampler	HiRISEGeoSampler	A TorchGeo geo-sampler that yields GeoSlice indices.	required
resolution	int	Output spatial resolution in pixels (square crop).	512
dtm_normalization	Literal['relative', 'log', 'linear']	"relative" (default) — local centering + fixed physical scale derived from centered_p98 in the stats file. "log" / "linear" — global quantile modes (kept for reference).	'relative'
random_flip	bool	Apply random horizontal/vertical flips + 90° rotations.	True
brightness_jitter	float	Max relative brightness perturbation on image only.	0.1
stats_path	str | None	Path to dataset_stats.json for normalization quantiles. Falls back to hardcoded Olympus-region defaults if None.	None
use_manifest	bool	If True, build / load a parquet manifest of clean patches so __getitem__ is O(1) and never retries.	True

Source code in src/depth_fm/data/adapter.py

def __init__(
        self,
        base_dataset: MarsHiRISEDTM,
        sampler: HiRISEGeoSampler,
        resolution: int = 512,
        dtm_normalization: Literal["relative", "log", "linear"] = "relative",
        random_flip: bool = True,
        random_jitter: bool = False,
        brightness_jitter: float = 0.1,
        stats_path: str | None = None,
        clip: bool = False,
        use_manifest: bool = True,
        manifest_workers: int = 16,
        manifest_dir: str | None = None,
        erode_radius: int = 2,
        multiprocessing_context="fork",
):
    super().__init__()
    self.multiprocessing_context = multiprocessing_context
    self.erode_radius = erode_radius
    self.base = base_dataset
    self.sampler = sampler
    self.resolution = resolution
    self.dtm_norm = dtm_normalization
    self.flip = random_flip
    self.random_jitter = random_jitter
    self.bright_jitter = brightness_jitter
    self.is_train = random_flip

    self.use_manifest = use_manifest
    self.manifest_workers = manifest_workers
    from cache import manifest_cache_dir
    self.manifest_dir = (
        Path(manifest_dir) if manifest_dir is not None
        else manifest_cache_dir(self.base.root)
    )
    self.clip = clip

    # Load global quantiles
    (
        self.elev_p02, self.elev_p98,
        self.img_p02, self.img_p98,
        self.elev_scale,
    ) = _load_quantiles(stats_path)

    # TODO: per-strip ortho normalization would be more correct than
    # per-patch, but the infrastructure doesn't yet thread strip identity
    # through the sampler.
    self.ortho_normalizer = LocalStripOrthoNormalizer(self.img_p02, self.img_p98)
    self.evel_normalizer = GlobalLogNormalizer(self.elev_scale, clip=clip)

    self._raw_indices = list(sampler)

    if self.use_manifest:
        self._init_manifest()
    else:
        self.clean_records = None
        logger.warning("Manifest disabled. Training will be slow due to on-the-fly validation.")

erode_valid_mask

erode_valid_mask(valid_mask: Tensor, erode_radius: int = 1) -> torch.Tensor

Erode a binary mask to trim noisy boundary pixels.

Source code in src/depth_fm/data/image_processing/mask_ops.py

def erode_valid_mask(valid_mask: torch.Tensor, erode_radius: int = 1) -> torch.Tensor:
    """Erode a binary mask to trim noisy boundary pixels."""
    if erode_radius <= 0:
        return valid_mask

    is_3d = valid_mask.ndim == 3
    if is_3d:
        valid_mask = valid_mask.unsqueeze(0)

    kernel_size = 2 * erode_radius + 1
    padded_mask = F.pad(
        valid_mask,
        pad=(erode_radius, erode_radius, erode_radius, erode_radius),
        mode="constant",
        value=1.0,
    )
    eroded_mask = -F.max_pool2d(
        -padded_mask, kernel_size=kernel_size, stride=1, padding=0
    )
    eroded_mask = (eroded_mask > 0.5).float()

    if is_3d:
        return eroded_mask.squeeze(0)
    return eroded_mask

fill_dtm_smart_diffusion

fill_dtm_smart_diffusion(tensor: Tensor, valid_mask: Tensor, iterations: int = 64, erode_radius: int = 2) -> torch.Tensor

Fill invalid regions using Laplacian diffusion (heat equation).

Source code in src/depth_fm/data/image_processing/void_filling.py

def fill_dtm_smart_diffusion(
        tensor: torch.Tensor,
        valid_mask: torch.Tensor,
        iterations: int = 64,
        erode_radius: int = 2,
) -> torch.Tensor:
    """Fill invalid regions using Laplacian diffusion (heat equation)."""
    is_3d = tensor.ndim == 3
    if is_3d:
        tensor = tensor.unsqueeze(0)
        valid_mask = valid_mask.unsqueeze(0)

    B, C, H, W = tensor.shape
    device = tensor.device
    dtype = tensor.dtype

    tensor = torch.nan_to_num(tensor, nan=0.0)

    valid_bool = erode_valid_mask(valid_mask, erode_radius) > 0.5

    valid_sum = (tensor * valid_bool.to(dtype)).sum(dim=[-2, -1], keepdim=True)
    valid_count = valid_bool.sum(dim=[-2, -1], keepdim=True).clamp(min=1.0)
    global_mean = valid_sum / valid_count

    filled = torch.where(valid_bool, tensor, global_mean)

    kernel = torch.ones((C, 1, 3, 3), device=device, dtype=dtype) / 9.0

    for _ in range(iterations):
        # Replicate padding stops physical image edges from dragging boundary
        # values down to 0.0 during the blur phase.
        padded_filled = F.pad(filled, pad=(1, 1, 1, 1), mode="replicate")
        blurred = F.conv2d(padded_filled, kernel, padding=0, groups=C)
        filled = torch.where(valid_mask, tensor, blurred)

    if is_3d:
        return filled.squeeze(0)
    return filled

fill_invalid_nearest_neighbor

fill_invalid_nearest_neighbor(tensor: Tensor, valid_mask: Tensor) -> torch.Tensor

Fill invalid regions using nearest-neighbor propagation via EDT.

Supports (H, W), (C, H, W), and (B, C, H, W) tensors.

Source code in src/depth_fm/data/image_processing/void_filling.py

def fill_invalid_nearest_neighbor(
        tensor: torch.Tensor,
        valid_mask: torch.Tensor
) -> torch.Tensor:
    """Fill invalid regions using nearest-neighbor propagation via EDT.

    Supports (H, W), (C, H, W), and (B, C, H, W) tensors.
    """
    device = tensor.device
    dtype = tensor.dtype

    arr_np = tensor.cpu().numpy()
    mask_np = valid_mask.squeeze().cpu().numpy() > 0

    if mask_np.all() or not mask_np.any():
        return tensor

    indices = ndimage.distance_transform_edt(
        ~mask_np,
        return_distances=False,
        return_indices=True,
    )
    iy, ix = indices

    if arr_np.ndim == 2:
        filled_np = arr_np[iy, ix]

    elif arr_np.ndim == 3:
        C = arr_np.shape[0]
        c_idx = np.arange(C)[:, None, None]
        iy_b = iy[None, :, :]
        ix_b = ix[None, :, :]
        filled_np = arr_np[c_idx, iy_b, ix_b]

    elif arr_np.ndim == 4:
        B, C = arr_np.shape[:2]
        b_idx = np.arange(B)[:, None, None, None]
        c_idx = np.arange(C)[None, :, None, None]
        iy_b = iy[None, None, :, :]
        ix_b = ix[None, None, :, :]
        filled_np = arr_np[b_idx, c_idx, iy_b, ix_b]

    else:
        raise ValueError(f"Unsupported tensor dimension: {arr_np.ndim}")

    return torch.from_numpy(filled_np).to(device=device, dtype=dtype)

fill_voids_gmrf

fill_voids_gmrf(image: Tensor, dtm: Tensor, valid_mask: Tensor, *, erode_radius: int = 2, connectivity: int = 4, tau: float = 1.0, nugget: float = 1e-06) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Fill voids in an orthoimage and DTM using GMRF conditional distribution.

One sparse linear solve per channel. No iteration. O(n^{3/2}) on 2D grids.

Returns (filled_image, filled_dtm, eroded_mask). The eroded mask records which pixels were considered trustworthy (1) vs infilled (0) — all pixels are valid in the filled outputs.

Source code in src/depth_fm/data/image_processing/void_filling.py

def fill_voids_gmrf(
        image: torch.Tensor,
        dtm: torch.Tensor,
        valid_mask: torch.Tensor,
        *,
        erode_radius: int = 2,
        connectivity: int = 4,
        tau: float = 1.0,
        nugget: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Fill voids in an orthoimage and DTM using GMRF conditional distribution.

    One sparse linear solve per channel. No iteration. O(n^{3/2}) on 2D grids.

    Returns (filled_image, filled_dtm, eroded_mask). The eroded mask records
    which pixels were considered trustworthy (1) vs infilled (0) — all pixels
    are valid in the filled outputs.
    """
    img_3d = image.ndim == 3
    dtm_3d = dtm.ndim == 3
    msk_3d = valid_mask.ndim == 3

    if img_3d:
        image = image.unsqueeze(0)
    if dtm_3d:
        dtm = dtm.unsqueeze(0)
    if msk_3d:
        valid_mask = valid_mask.unsqueeze(0)

    device = image.device
    img_dtype = image.dtype
    dtm_dtype = dtm.dtype

    eroded = erode_valid_mask(valid_mask, erode_radius)
    valid = eroded[0, 0].cpu().numpy() > 0.5

    img_np = np.nan_to_num(image[0].cpu().numpy(), nan=0.0).astype(np.float64)
    dtm_np = np.nan_to_num(dtm[0].cpu().numpy(), nan=0.0).astype(np.float64)

    if not (~valid).any():
        eroded_out = eroded
        if img_3d:
            image = image.squeeze(0)
        if dtm_3d:
            dtm = dtm.squeeze(0)
        if msk_3d:
            eroded_out = eroded_out.squeeze(0)
        return image, dtm, eroded_out

    H, W = valid.shape

    L = _build_grid_laplacian(H, W, connectivity)
    Q = tau * L

    obs_idx = np.where(valid.ravel())[0]
    void_idx = np.where(~valid.ravel())[0]

    logger.debug(
        "GMRF fill: %d void pixels (%.1f%%), %d-connected, H=%d W=%d",
        len(void_idx), 100.0 * len(void_idx) / (H * W), connectivity, H, W,
    )

    filled_dtm = dtm_np.copy()
    filled_img = img_np.copy()

    for c in range(dtm_np.shape[0]):
        filled_dtm[c] = _gmrf_fill_channel(dtm_np[c], void_idx, obs_idx, Q, nugget)

    for c in range(img_np.shape[0]):
        filled_img[c] = _gmrf_fill_channel(img_np[c], void_idx, obs_idx, Q, nugget)

    filled_img_t = torch.from_numpy(filled_img).unsqueeze(0).to(device=device, dtype=img_dtype)
    filled_dtm_t = torch.from_numpy(filled_dtm).unsqueeze(0).to(device=device, dtype=dtm_dtype)

    if img_3d:
        filled_img_t = filled_img_t.squeeze(0)
    if dtm_3d:
        filled_dtm_t = filled_dtm_t.squeeze(0)
    if msk_3d:
        eroded = eroded.squeeze(0)

    return filled_img_t, filled_dtm_t, eroded

fill_voids_kriging

fill_voids_kriging(image: Tensor, dtm: Tensor, valid_mask: Tensor, *, erode_radius: int = 2, max_training_points: int = 2500, variogram_model: str = 'linear', seed: int = 42) -> tuple[torch.Tensor, torch.Tensor]

Fill voids in an orthoimage and DTM with a single kriging pass per channel.

Optimised for 512x512 tiles: subsamples valid pixels for training, then predicts every void pixel at once. No iteration, no per-void labelling, no diffusion loop.

Source code in src/depth_fm/data/image_processing/void_filling.py

def fill_voids_kriging(
        image: torch.Tensor,
        dtm: torch.Tensor,
        valid_mask: torch.Tensor,
        *,
        erode_radius: int = 2,
        max_training_points: int = 2500,
        variogram_model: str = "linear",
        seed: int = 42,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fill voids in an orthoimage and DTM with a single kriging pass per channel.

    Optimised for 512x512 tiles: subsamples valid pixels for training, then
    predicts every void pixel at once. No iteration, no per-void labelling,
    no diffusion loop.
    """
    rng = np.random.default_rng(seed)

    img_3d = image.ndim == 3
    dtm_3d = dtm.ndim == 3
    msk_3d = valid_mask.ndim == 3

    if img_3d:
        image = image.unsqueeze(0)
    if dtm_3d:
        dtm = dtm.unsqueeze(0)
    if msk_3d:
        valid_mask = valid_mask.unsqueeze(0)

    device = image.device
    img_dtype = image.dtype
    dtm_dtype = dtm.dtype

    eroded = erode_valid_mask(valid_mask, erode_radius)
    valid = eroded[0, 0].cpu().numpy() > 0.5
    void = ~valid

    img_np = np.nan_to_num(image[0].cpu().numpy(), nan=0.0).astype(np.float64)
    dtm_np = np.nan_to_num(dtm[0].cpu().numpy(), nan=0.0).astype(np.float64)

    if not void.any():
        if img_3d:
            image = image.squeeze(0)
        if dtm_3d:
            dtm = dtm.squeeze(0)
        return image, dtm

    y_valid, x_valid = np.where(valid)
    y_void, x_void = np.where(void)

    n_valid = len(y_valid)
    if n_valid < 6:
        logger.warning("Fewer than 6 valid pixels — returning inputs unchanged.")
        if img_3d:
            image = image.squeeze(0)
        if dtm_3d:
            dtm = dtm.squeeze(0)
        return image, dtm

    if n_valid > max_training_points:
        idx = rng.choice(n_valid, size=max_training_points, replace=False)
        y_train = y_valid[idx].astype(np.float64)
        x_train = x_valid[idx].astype(np.float64)
    else:
        y_train = y_valid.astype(np.float64)
        x_train = x_valid.astype(np.float64)

    y_pred = y_void.astype(np.float64)
    x_pred = x_void.astype(np.float64)

    filled_img = img_np.copy()
    filled_dtm = dtm_np.copy()

    for c in range(dtm_np.shape[0]):
        z_train = dtm_np[c, y_train.astype(int), x_train.astype(int)]
        try:
            uk = UniversalKriging(
                x_train, y_train, z_train,
                variogram_model=variogram_model,
                drift_terms=["regional_linear"],
                verbose=False,
                enable_plotting=False,
            )
            z_pred, _ = uk.execute("points", x_pred, y_pred)
            filled_dtm[c, y_void, x_void] = np.asarray(z_pred).ravel()
        except Exception as e:
            logger.warning("DTM kriging failed (ch %d): %s", c, e)

    for c in range(img_np.shape[0]):
        z_train = img_np[c, y_train.astype(int), x_train.astype(int)]
        try:
            ok = OrdinaryKriging(
                x_train, y_train, z_train,
                variogram_model=variogram_model,
                verbose=False,
                enable_plotting=False,
            )
            z_pred, _ = ok.execute("points", x_pred, y_pred)
            filled_img[c, y_void, x_void] = np.asarray(z_pred).ravel()
        except Exception as e:
            logger.warning("Image kriging failed (ch %d): %s", c, e)

    filled_img_t = torch.from_numpy(filled_img).unsqueeze(0).to(device=device, dtype=img_dtype)
    filled_dtm_t = torch.from_numpy(filled_dtm).unsqueeze(0).to(device=device, dtype=dtm_dtype)

    if img_3d:
        filled_img_t = filled_img_t.squeeze(0)
    if dtm_3d:
        filled_dtm_t = filled_dtm_t.squeeze(0)

    return filled_img_t, filled_dtm_t

compute_artifact_multipliers

compute_artifact_multipliers(seam_heatmap: ndarray, valid_mask: ndarray, threshold: float = 0.2) -> tuple[float, float, np.ndarray, np.ndarray]

Compute structural multipliers that distinguish seams from natural features.

Returns (span_ratio, sparsity, hot_mask, labels): * span_ratio — defeats short craters; true seams cross the whole tile. * sparsity — defeats dense dunes; true seams are a singular line.

Source code in src/depth_fm/data/image_processing/seam_detection.py

def compute_artifact_multipliers(
        seam_heatmap: np.ndarray,
        valid_mask: np.ndarray,
        threshold: float = 0.2,
) -> tuple[float, float, np.ndarray, np.ndarray]:
    """Compute structural multipliers that distinguish seams from natural features.

    Returns (span_ratio, sparsity, hot_mask, labels):
    * span_ratio — defeats short craters; true seams cross the whole tile.
    * sparsity   — defeats dense dunes; true seams are a singular line.
    """
    max_val = np.max(seam_heatmap)
    if max_val <= 0:
        return 0.0, 1.0

    hot_mask = (seam_heatmap > max_val * threshold)

    # SPARSITY (defeats repeating textures like dunes)
    hot_area = np.count_nonzero(hot_mask & valid_mask)
    valid_area = max(np.count_nonzero(valid_mask), 1)
    density = hot_area / valid_area
    # True seams cover ~2-4% of the image; dune fields cover >15%.
    sparsity = math.exp(-density * 15.0)

    # STRUCTURAL SPAN (defeats short, isolated craters)
    binary = (hot_mask.astype(np.uint8)) * 255

    kernel = np.ones((9, 9), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)

    span_ratio = 0.0
    if num_labels > 1:
        H, W = seam_heatmap.shape
        max_span = 0.0
        for i in range(1, num_labels):
            w = stats[i, cv2.CC_STAT_WIDTH]
            h = stats[i, cv2.CC_STAT_HEIGHT]
            span = math.hypot(w, h)
            if span > max_span:
                max_span = span
        max_possible_span = float(max(H, W))
        span_ratio = min(max_span / max_possible_span, 1.0)

    return float(span_ratio), float(sparsity), hot_mask, labels

compute_piecewise_linearity

compute_piecewise_linearity(seam_heatmap: ndarray, threshold_ratio: float = 0.3) -> tuple[float, np.ndarray]

Score how piecewise-linear the high-scoring pixels are (handles corners).

Returns (linearity_score, line_mask) where linearity is in [0.0, 1.0]: 1.0 = highly structured/linear, 0.0 = curved/messy.

Source code in src/depth_fm/data/image_processing/seam_detection.py

def compute_piecewise_linearity(seam_heatmap: np.ndarray, threshold_ratio: float = 0.3) -> tuple[float, np.ndarray]:
    """Score how piecewise-linear the high-scoring pixels are (handles corners).

    Returns (linearity_score, line_mask) where linearity is in [0.0, 1.0]:
    1.0 = highly structured/linear, 0.0 = curved/messy.
    """
    heatmap_norm = cv2.normalize(seam_heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    _, binary_map = cv2.threshold(heatmap_norm, int(255 * threshold_ratio), 255, cv2.THRESH_BINARY)

    # Bridge missing line segments before edge detection so fragmented seams
    # heal into a continuous solid line. 15×15 closing merges blobs that are
    # separated by up to ~15 pixels.
    close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    closed_map = cv2.morphologyEx(binary_map, cv2.MORPH_CLOSE, close_kernel)

    edges = cv2.Canny(closed_map, 50, 150, apertureSize=3)

    lines = cv2.HoughLinesP(edges, rho=2, theta=np.pi / 180, threshold=50,
                            minLineLength=40, maxLineGap=30)

    line_mask = np.zeros_like(binary_map)

    if lines is None:
        return 0.0, line_mask

    for line in lines:
        x1, y1, x2, y2 = line[0]
        cv2.line(line_mask, (x1, y1), (x2, y2), 255, thickness=4)

    valid_signal_pixels = np.count_nonzero(binary_map)
    if valid_signal_pixels == 0:
        return 0.0, line_mask

    structured_pixels = np.count_nonzero(cv2.bitwise_and(binary_map, line_mask))

    # True seams (even with corners) score close to 1.0; natural curves fall
    # apart in the Hough transform and score low.
    return structured_pixels / valid_signal_pixels, line_mask

compute_spatial_isolation

compute_spatial_isolation(score_map: Tensor, valid_mask: Tensor, best_x: int, best_y: int, angle_rad: float, profile_length: int = 50, exclusion_zone: int = 12) -> tuple[float, Optional[tuple]]

Sample a perpendicular slice across the seam.

Returns the ratio of the central peak to the surrounding parallel background, plus the raw profile data arrays for visualization.

Source code in src/depth_fm/data/image_processing/seam_detection.py

def compute_spatial_isolation(
        score_map: torch.Tensor,
        valid_mask: torch.Tensor,
        best_x: int,
        best_y: int,
        angle_rad: float,
        profile_length: int = 50,
        exclusion_zone: int = 12,
) -> tuple[float, Optional[tuple]]:
    """Sample a perpendicular slice across the seam.

    Returns the ratio of the central peak to the surrounding parallel
    background, plus the raw profile data arrays for visualization.
    """
    H, W = score_map.shape
    device = score_map.device

    px = -math.sin(angle_rad)
    py = math.cos(angle_rad)

    t = torch.arange(-profile_length, profile_length + 1, device=device, dtype=torch.float32)
    grid_x = best_x + t * px
    grid_y = best_y + t * py

    in_bounds = (grid_x >= 0) & (grid_x < W) & (grid_y >= 0) & (grid_y < H)
    t = t[in_bounds]
    grid_x = grid_x[in_bounds]
    grid_y = grid_y[in_bounds]

    if len(t) == 0:
        return 0.0, None

    norm_x = (grid_x / (W - 1)) * 2 - 1
    norm_y = (grid_y / (H - 1)) * 2 - 1
    grid = torch.stack([norm_x, norm_y], dim=-1).view(1, 1, -1, 2)

    score_map_4d = score_map.unsqueeze(0).unsqueeze(0)
    valid_mask_4d = valid_mask.unsqueeze(0).unsqueeze(0)

    samples = F.grid_sample(score_map_4d, grid, mode="bilinear", align_corners=True).squeeze()
    v_samples = F.grid_sample(valid_mask_4d, grid, mode="nearest", align_corners=True).squeeze()

    center_mask = (t.abs() <= exclusion_zone) & (v_samples > 0)
    bg_mask = (t.abs() > exclusion_zone) & (v_samples > 0)

    profile_data = (
        t.cpu().numpy(),
        samples.cpu().numpy(),
        center_mask.cpu().numpy(),
        bg_mask.cpu().numpy(),
    )

    if not center_mask.any():
        return 0.0, profile_data
    if not bg_mask.any():
        # No valid background — treat as an isolated edge.
        return 5.0, profile_data

    center_max = samples[center_mask].amax().item()
    bg_mean = samples[bg_mask].mean().item()

    if bg_mean < 1e-6:
        return 20.0, profile_data

    return center_max / bg_mean, profile_data

detect_seam_artifact

detect_seam_artifact(ortho: Tensor, elevation: Tensor, valid_mask: Tensor, line_length: int = 41, num_angles: int = 12, side_offset: int = 2, min_valid_ratio: float = 0.6, ortho_weight: float = 1.0, dtm_weight: float = 0.3, erosion_kernel: int = 9, seam_threshold: float = 2.4, return_diagnostics: bool = False) -> SeamResult

Detect seam artifacts (mosaicking discontinuities) in a HiRISE patch.

Returns a SeamResult. With return_diagnostics=True, the result also carries seam_heatmap, cohens_d_heatmap, and per_angle_max, used by the visualization and refinement UI to show where and along which angle the detector fired.

Source code in src/depth_fm/data/image_processing/seam_detection.py

@torch.no_grad()
def detect_seam_artifact(
        ortho: torch.Tensor,
        elevation: torch.Tensor,
        valid_mask: torch.Tensor,
        line_length: int = 41,
        num_angles: int = 12,
        side_offset: int = 2,
        min_valid_ratio: float = 0.6,
        ortho_weight: float = 1.0,
        dtm_weight: float = 0.3,
        erosion_kernel: int = 9,
        seam_threshold: float = 2.4,
        return_diagnostics: bool = False,
) -> SeamResult:
    """Detect seam artifacts (mosaicking discontinuities) in a HiRISE patch.

    Returns a `SeamResult`. With `return_diagnostics=True`, the result also
    carries `seam_heatmap`, `cohens_d_heatmap`, and `per_angle_max`, used by
    the visualization and refinement UI to show *where* and *along which
    angle* the detector fired.
    """
    ortho = _ensure_bchw(ortho).float()
    elevation = _ensure_bchw(elevation).float()
    valid_mask = _ensure_bchw(valid_mask).bool()
    device, dtype = ortho.device, ortho.dtype
    ortho_gray = ortho.mean(dim=1, keepdim=True)

    H, W = ortho_gray.shape[-2:]

    ortho_grad = _sharp_grad_mag(ortho_gray, valid_mask)
    dtm_grad = _sharp_grad_mag(elevation, valid_mask)

    pad_e = erosion_kernel // 2
    invalid_f = (~valid_mask).float()
    dilated = F.max_pool2d(invalid_f, kernel_size=erosion_kernel, stride=1, padding=pad_e)
    eroded_valid = (dilated == 0.0)
    ev_f = eroded_valid.float()

    empty_result = SeamResult(
        seam_score=0.0, ortho_score=0.0, dtm_score=0.0, cohens_d=0.0,
        best_angle_rad=0.0, best_y=H // 2, best_x=W // 2,
        line_length=line_length, num_angles=num_angles,
        composite_score=0, sparsity=0, span=0,
        is_seam=False,
    )

    if ev_f.sum() < 100:
        return empty_result

    def _bg(signal):
        num = (signal * ev_f).sum()
        den = ev_f.sum().clamp(min=1)
        return (num / den).clamp(min=1e-6)

    bg_ortho, bg_dtm = _bg(ortho_grad), _bg(dtm_grad)

    line_ker, left_ker, right_ker, pad = _build_oriented_kernels(
        line_length, num_angles, side_offset, device, dtype
    )

    def _line_avg(signal, ker):
        num = F.conv2d(signal * ev_f, ker, padding=pad)
        cnt = F.conv2d(ev_f, ker, padding=pad)
        return num / cnt.clamp(min=1.0), cnt

    ortho_line_avg, line_cnt = _line_avg(ortho_grad, line_ker)
    dtm_line_avg, _ = _line_avg(dtm_grad, line_ker)

    def _moments(signal, ker):
        n = F.conv2d(ev_f, ker, padding=pad).clamp(min=1.0)
        s = F.conv2d(signal * ev_f, ker, padding=pad)
        s2 = F.conv2d((signal ** 2) * ev_f, ker, padding=pad)
        mean = s / n
        var = (s2 / n - mean ** 2).clamp(min=0.0)
        return mean, var, n

    mu_l, var_l, n_l = _moments(ortho_gray, left_ker)
    mu_r, var_r, n_r = _moments(ortho_gray, right_ker)
    global_std = ortho_gray[valid_mask].std().clamp(min=1e-4)
    std_floor = 0.1 * global_std
    pooled_std = torch.sqrt(((var_l + var_r) / 2.0).clamp(min=0.0)) + std_floor
    cohens_d = (mu_l - mu_r).abs() / pooled_std

    mu_l_dtm, var_l_dtm, _ = _moments(elevation, left_ker)
    mu_r_dtm, var_r_dtm, _ = _moments(elevation, right_ker)

    global_dtm_std = elevation[valid_mask].std().clamp(min=1e-4)
    dtm_std_floor = 0.1 * global_dtm_std
    pooled_dtm_std = torch.sqrt(((var_l_dtm + var_r_dtm) / 2.0).clamp(min=0.0)) + dtm_std_floor

    dtm_cohens_d = (mu_l_dtm - mu_r_dtm).abs() / pooled_dtm_std

    min_line = min_valid_ratio * line_length
    min_side = min_valid_ratio * line_length * 0.5
    valid_q = (line_cnt >= min_line) & (n_l >= min_side) & (n_r >= min_side)
    if not valid_q.any():
        return empty_result

    ortho_norm = ortho_line_avg / bg_ortho
    dtm_norm = dtm_line_avg / bg_dtm
    gradient_term = ortho_weight * ortho_norm + dtm_weight * dtm_norm

    # Trigger if EITHER ortho OR dtm has a massive distribution shift.
    distribution_gate = torch.clamp(torch.maximum(cohens_d, dtm_cohens_d), min=0.1, max=3.0)

    combined = gradient_term * distribution_gate

    combined_masked = torch.where(valid_q, combined, torch.full_like(combined, -1e10))

    cm0 = combined_masked[0]
    flat = cm0.flatten().argmax()
    a_idx = int(flat // (H * W))
    rem = int(flat % (H * W))
    y_idx = rem // W
    x_idx = rem % W
    best_angle_rad = math.pi * a_idx / num_angles

    seam_score = float(cm0.amax().item())

    heatmap = cm0.amax(dim=0).cpu().numpy()
    heatmap[heatmap < 0] = 0.0

    valid_np = ev_f[0, 0].cpu().numpy().astype(bool)

    threshold = 0.25

    span, sparsity, diag_hot_mask, diag_labels = compute_artifact_multipliers(heatmap, valid_np, threshold)

    linearity, diag_hough_lines = compute_piecewise_linearity(heatmap, threshold_ratio=threshold)

    isolation_score, diag_iso_profile = compute_spatial_isolation(
        score_map=cm0.amax(dim=0),
        valid_mask=ev_f[0, 0],
        best_x=int(x_idx),
        best_y=int(y_idx),
        angle_rad=best_angle_rad,
    )

    isolation_mult = min(max((isolation_score - 1.5) / 1.5, 0.0), 1.0)
    composite_score = seam_score * (0.5 + 0.5 * span) * sparsity * linearity * isolation_mult

    result = SeamResult(
        seam_score=seam_score,
        ortho_score=float(torch.where(valid_q, ortho_norm, torch.full_like(ortho_norm, -1e10)).amax().item()),
        dtm_score=float(torch.where(valid_q, dtm_norm, torch.full_like(dtm_norm, -1e10)).amax().item()),
        cohens_d=float(torch.where(valid_q, cohens_d, torch.zeros_like(cohens_d)).amax().item()),
        best_angle_rad=best_angle_rad,
        best_y=int(y_idx),
        best_x=int(x_idx),
        line_length=line_length,
        span=span,
        sparsity=sparsity,
        num_angles=num_angles,
        composite_score=composite_score,
        is_seam=composite_score > seam_threshold,
    )

    if return_diagnostics:
        d0 = torch.where(valid_q, cohens_d, torch.full_like(cohens_d, float("nan")))[0]
        d_heatmap = d0.amax(dim=0).cpu().numpy()
        per_angle = cm0.amax(dim=(1, 2)).cpu().numpy()
        per_angle[per_angle < 0] = 0.0

        result.seam_heatmap = heatmap
        result.cohens_d_heatmap = d_heatmap
        result.per_angle_max = per_angle

        result.diag_hot_mask = diag_hot_mask
        result.diag_closed_components = diag_labels
        result.diag_hough_lines = diag_hough_lines
        result.diag_isolation_profile = diag_iso_profile

    return result

is_tin_artifact

is_tin_artifact(elevation: Tensor, valid_mask: Tensor, kernel_size: int = 32) -> float

Scale-invariant TIN-artifact detection using localized maximum density.

Returns the maximum local fraction of zero-curvature pixels within any kernel_size × kernel_size window. Patches with stretched-triangle TIN artifacts produce values near 1.0; natural terrain produces values <0.5.

Source code in src/depth_fm/data/image_processing/seam_detection.py

def is_tin_artifact(
        elevation: torch.Tensor,
        valid_mask: torch.Tensor,
        kernel_size: int = 32,
) -> float:
    """Scale-invariant TIN-artifact detection using localized maximum density.

    Returns the maximum local fraction of zero-curvature pixels within any
    `kernel_size × kernel_size` window. Patches with stretched-triangle TIN
    artifacts produce values near 1.0; natural terrain produces values <0.5.
    """
    if elevation.dim() == 2:
        elevation = elevation.view(1, 1, elevation.shape[0], elevation.shape[1])
        valid_mask = valid_mask.view(1, 1, valid_mask.shape[0], valid_mask.shape[1])

    safe_elev = elevation.clone()
    safe_elev[~valid_mask] = 0.0

    laplacian_kernel = torch.tensor([[[[0.0, 1.0, 0.0],
                                       [1.0, -4.0, 1.0],
                                       [0.0, 1.0, 0.0]]]], device=elevation.device)
    laplacian = F.conv2d(safe_elev, laplacian_kernel, padding=1)

    invalid_mask = (~valid_mask).float()
    dilated_invalid = F.max_pool2d(invalid_mask, kernel_size=3, stride=1, padding=1)
    eroded_valid = (dilated_invalid == 0.0).float()

    # True (1.0) if curvature ≈ 0 AND the pixel is deeply valid.
    zero_curvature_mask = ((laplacian.abs() < 1e-2) * eroded_valid.bool()).float()

    # Slide kernel_size×kernel_size window across every pixel via avg_pool2d.
    local_planar_sum = F.avg_pool2d(zero_curvature_mask, kernel_size=kernel_size, stride=1)
    local_valid_sum = F.avg_pool2d(eroded_valid, kernel_size=kernel_size, stride=1)

    safe_valid_sum = torch.clamp(local_valid_sum, min=1e-6)
    local_density = local_planar_sum / safe_valid_sum

    # Require ≥50% valid data in the window for the statistic to be sound.
    valid_window_mask = local_valid_sum >= 0.5

    if not valid_window_mask.any():
        return False

    return local_density[valid_window_mask].max().item()

estimate_sun_vector_irls

estimate_sun_vector_irls(dtm: Tensor, ortho: Tensor, valid_mask: Tensor, max_iter: int = 15, tol: float = 0.0001) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Decoupled IRLS sun-vector estimator.

Separates ambient light estimation from the linear system to prevent the Nz / bias collinearity trap from inverting the sun vector. Returns (sun_vec, intensity, ambient).

Source code in src/depth_fm/data/image_processing/sun_vector.py

@torch.no_grad()
def estimate_sun_vector_irls(
        dtm: torch.Tensor,
        ortho: torch.Tensor,
        valid_mask: torch.Tensor,
        max_iter: int = 15,
        tol: float = 1e-4,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Decoupled IRLS sun-vector estimator.

    Separates ambient light estimation from the linear system to prevent the
    Nz / bias collinearity trap from inverting the sun vector. Returns
    (sun_vec, intensity, ambient).
    """
    device = dtm.device

    def get_defaults():
        default_sun = F.normalize(torch.tensor([0.5, -0.5, 1.0], device=device), p=2, dim=0)
        return default_sun, torch.tensor(1.0, device=device), torch.tensor(0.05, device=device)

    if dtm.ndim == 4:
        dtm = dtm[0]
    if ortho.ndim == 4:
        ortho = ortho[0]
    if valid_mask.ndim == 4:
        valid_mask = valid_mask[0]
    if ortho.shape[0] == 3:
        ortho = ortho.mean(dim=0, keepdim=True)

    sobel_x = torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]], device=device) / 8.0
    sobel_y = torch.tensor([[-1., -2., -1.], [0., 0., 0.], [1., 2., 1.]], device=device) / 8.0
    spatial_scale = max(dtm.shape[-2], dtm.shape[-1]) / 2.0
    padded_dtm = F.pad(dtm.unsqueeze(0), (1, 1, 1, 1), mode="replicate")

    n_x = -F.conv2d(padded_dtm, sobel_x.view(1, 1, 3, 3)) * spatial_scale
    n_y = -F.conv2d(padded_dtm, sobel_y.view(1, 1, 3, 3)) * spatial_scale
    n_z = torch.ones_like(n_x)

    normals = F.normalize(torch.cat([n_x, n_y, n_z], dim=1), p=2, dim=1).squeeze(0)

    mask = valid_mask.squeeze(0).bool()
    zero_mask = ortho.squeeze(0) > 1e-4
    final_mask = mask & zero_mask

    N_flat = normals[:, final_mask].t()
    Y_raw = ortho.squeeze(0)[final_mask].unsqueeze(1)

    if N_flat.shape[0] < 100:
        return get_defaults()

    # Decoupled ambient estimation: 1st percentile is a robust proxy for the
    # secondary scattering / ambient floor in deep shadows.
    ambient_est = torch.quantile(Y_raw, 0.01)
    Y_flat = torch.clamp(Y_raw - ambient_est, min=0.0)

    H = N_flat
    beta = torch.linalg.lstsq(H, Y_flat).solution
    c = 1.345

    for _ in range(max_iter):
        residuals = Y_flat - torch.mm(H, beta)
        median_res = torch.median(residuals)
        mad = torch.median(torch.abs(residuals - median_res))
        sigma = (mad / 0.67449) + 1e-6

        r_stand = torch.abs(residuals / sigma)
        weights = torch.clamp(c / (r_stand + 1e-8), max=1.0)

        w_sqrt = torch.sqrt(weights)
        H_w = H * w_sqrt
        Y_w = Y_flat * w_sqrt

        beta_new = torch.linalg.lstsq(H_w, Y_w).solution

        change = torch.norm(beta_new - beta, p=2)
        # Every successful solve becomes the next IRLS iterate, including
        # the final update when the iteration budget is exhausted.
        beta = beta_new
        if change < tol:
            break

    k = beta[:3, 0]
    # Failsafe: with the bias stripped, if anomalous geometry still pushes Nz
    # negative we reflect across the horizon line to maintain physical validity.
    if k[2] < 0:
        k[2] = -k[2]

    intensity = torch.norm(k, p=2).clamp(min=1e-4)
    sun_vec = F.normalize(k, p=2, dim=0)

    return sun_vec, intensity, ambient_est

estimate_sun_vector_ols

estimate_sun_vector_ols(dtm: Tensor, ortho: Tensor, valid_mask: Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Estimates the sun vector [sx, sy, sz] using Ordinary Least Squares.

Handles shapes (C, H, W) or (1, C, H, W). Batch size must be 1.

Source code in src/depth_fm/data/image_processing/sun_vector.py

@deprecated(
    "Use estimate_sun_vector_irls instead, this does not calculate the z sun vector correctly and can render it to have negative values."
)
def estimate_sun_vector_ols(
        dtm: torch.Tensor,
        ortho: torch.Tensor,
        valid_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Estimates the sun vector [sx, sy, sz] using Ordinary Least Squares.

    Handles shapes (C, H, W) or (1, C, H, W). Batch size must be 1.
    """
    device = dtm.device

    def get_defaults():
        default_sun = F.normalize(torch.tensor([0.5, -0.5, 1.0], device=device), p=2, dim=0)
        return default_sun, torch.tensor(1.0, device=device), torch.tensor(0.3, device=device)

    if dtm.ndim == 4:
        assert dtm.shape[0] == 1, (
            "The batch size should be 1. estimate_sun_vector_irls currently only works for a batch size of 1"
        )
        dtm = dtm[0]
    if ortho.ndim == 4:
        ortho = ortho[0]
    if valid_mask.ndim == 4:
        valid_mask = valid_mask[0]

    if ortho.shape[0] == 3:
        ortho = ortho.mean(dim=0, keepdim=True)

    sobel_x = torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]], device=device) / 8.0
    sobel_y = torch.tensor([[-1., -2., -1.], [0., 0., 0.], [1., 2., 1.]], device=device) / 8.0

    spatial_scale = max(dtm.shape[-2], dtm.shape[-1]) / 2.0
    padded_dtm = F.pad(dtm.unsqueeze(0), (1, 1, 1, 1), mode="replicate")

    n_x = -F.conv2d(padded_dtm, sobel_x.view(1, 1, 3, 3)) * spatial_scale
    n_y = -F.conv2d(padded_dtm, sobel_y.view(1, 1, 3, 3)) * spatial_scale
    n_z = torch.ones_like(n_x)

    normals = torch.cat([n_x, n_y, n_z], dim=1)
    normals = F.normalize(normals, p=2, dim=1).squeeze(0)

    mask = valid_mask.squeeze(0).bool()

    ortho_valid = ortho.squeeze(0)[mask]
    if len(ortho_valid) == 0:
        return get_defaults()

    intensity_threshold = torch.quantile(ortho_valid, 0.05)
    shadow_mask = ortho.squeeze(0) > intensity_threshold
    final_mask = mask & shadow_mask

    N_flat = normals[:, final_mask].t()
    Y_flat = ortho.squeeze(0)[final_mask].unsqueeze(1)

    if N_flat.shape[0] < 100:
        return get_defaults()

    ones = torch.ones((N_flat.shape[0], 1), device=device)
    A = torch.cat([N_flat, ones], dim=1)

    x = torch.linalg.lstsq(A, Y_flat).solution

    k = x[:3, 0]
    ambient = x[3, 0]
    intensity = torch.norm(k, p=2)
    sun_vec = F.normalize(k, p=2, dim=0)

    return sun_vec, intensity, ambient

compute_topographic_residual

compute_topographic_residual(elevation: Tensor, valid_mask: Tensor) -> float

Fit a 2D plane to elevation and return the RMS residual.

Removes macroscopic slopes so the returned scalar isolates true topographic roughness (used as a manifest filter to reject overly-flat patches).

Source code in src/depth_fm/data/image_processing/terrain.py

def compute_topographic_residual(elevation: torch.Tensor, valid_mask: torch.Tensor) -> float:
    """Fit a 2D plane to elevation and return the RMS residual.

    Removes macroscopic slopes so the returned scalar isolates true topographic
    roughness (used as a manifest filter to reject overly-flat patches).
    """
    if elevation.ndim > 2:
        elevation = elevation.squeeze()
        valid_mask = valid_mask.squeeze()

    H, W = elevation.shape

    y = torch.linspace(-1, 1, H, dtype=elevation.dtype, device=elevation.device)
    x = torch.linspace(-1, 1, W, dtype=elevation.dtype, device=elevation.device)
    Y, X = torch.meshgrid(y, x, indexing="ij")

    valid_bool = valid_mask.bool()
    if not valid_bool.any():
        return 0.0

    X_v = X[valid_bool].unsqueeze(1)
    Y_v = Y[valid_bool].unsqueeze(1)
    Z_v = elevation[valid_bool].unsqueeze(1)

    A = torch.cat([X_v, Y_v, torch.ones_like(X_v)], dim=1)
    w = torch.linalg.lstsq(A, Z_v).solution

    Z_pred = A @ w
    return torch.sqrt(torch.mean((Z_v - Z_pred) ** 2)).item()