XL2     &                cRڳJ     TypeV2Obj ID                DDirYTiD'k8+ҦEcAlgoEcMEcN EcBSize   EcIndexEcDistCSumAlgoPartNumsPartETags 85b8e43a254cdd5dbe15987fb89a5fafPartSizes8PartASizes8PartIdx Size8MTimecRڳJMetaSysx-minio-internal-inline-datatruex-rustfs-internal-inline-datatrueMetaUsretag 85b8e43a254cdd5dbe15987fb89a5fafcontent-typetext/x-python-scriptv m|ʌnullX5&$/jf~j?Wb#from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

from .calibration import apply_gain_map, generate_gain_map
from .ellipse_model import fit_ellipse_from_mask
from .io_utils import ensure_dir, read_gray_image, save_image
from .preprocess import build_ring_mask, denoise_image, extract_ring_band_mask
from .unwrap import enhance_unwrapped_strip, estimate_band_offsets, unwrap_ring_band
from .visualize import save_batch_montage


def natural_key(path):
    stem = Path(path).stem
    try:
        return int(stem)
    except ValueError:
        return stem


def build_theta_samples(n_theta=1880):
    return np.linspace(0.0, 2.0 * np.pi, num=n_theta, endpoint=False, dtype=np.float32)


def collapse_strip(strip, mode="max"):
    strip = np.asarray(strip, dtype=np.float32)
    if mode == "sum":
        vector = np.sum(strip, axis=0)
    else:
        vector = np.max(strip, axis=0)
    return vector.astype(np.float32)


def detrend_surface(surface, sigma_y=22):
    surface = surface.astype(np.float32)
    background = cv2.GaussianBlur(surface, (0, 0), sigmaX=float(sigma_y), sigmaY=0)
    return surface - background


def normalize_surface(surface):
    if surface.size == 0:
        return surface.astype(np.uint8)
    norm = cv2.normalize(surface, None, 0, 255, cv2.NORM_MINMAX)
    return norm.astype(np.uint8)


def enhance_surface(surface, clahe_clip=2.0, sigma_y=22):
    detrended = detrend_surface(surface, sigma_y=sigma_y)
    normalized = normalize_surface(detrended)
    clahe = cv2.createCLAHE(clipLimit=float(clahe_clip), tileGridSize=(8, 8))
    return clahe.apply(normalized)


@dataclass
class CharacterStitchResult:
    image_paths: list
    vectors: list
    raw_stack: np.ndarray
    stretched_stack: np.ndarray
    enhanced_stack: np.ndarray


def process_character_image(
    image_path,
    calibration=None,
    n_theta=1880,
    strip_height=96,
    collapse_mode="max",
):
    image = read_gray_image(image_path)
    ring_mask = build_ring_mask(denoise_image(image))
    band_mask = extract_ring_band_mask(image, ring_mask)
    ellipse = fit_ellipse_from_mask(ring_mask)
    theta_samples = build_theta_samples(n_theta)

    corrected = image
    if calibration is not None:
        gain_map = generate_gain_map(
            image.shape,
            ellipse,
            calibration.theta_samples,
            calibration.gains,
            blur_sigma=20,
        )
        corrected = apply_gain_map(image, gain_map, band_mask=band_mask)

    corrected_band_mask = extract_ring_band_mask(corrected, ring_mask)
    inner_offset, outer_offset = estimate_band_offsets(corrected_band_mask, ellipse, theta_samples)
    strip = unwrap_ring_band(
        corrected,
        ellipse,
        theta_samples,
        inner_offset,
        outer_offset,
        out_height=strip_height,
    )
    strip = enhance_unwrapped_strip(strip)
    vector = collapse_strip(strip, mode=collapse_mode)
    return {
        "image": corrected,
        "ring_mask": ring_mask,
        "band_mask": band_mask,
        "ellipse": ellipse,
        "strip": strip,
        "vector": vector,
    }


def stitch_character_batch(
    input_dir,
    output_dir,
    calibration=None,
    n_theta=1880,
    strip_height=96,
    target_height=400,
    collapse_mode="max",
):
    input_dir = Path(input_dir)
    output_dir = ensure_dir(output_dir)
    image_paths = sorted(input_dir.glob("*.bmp"), key=natural_key)
    if not image_paths:
        raise FileNotFoundError(f"No BMP files found in {input_dir}")

    strips = []
    vectors = []
    labels = []
    for image_path in image_paths:
        result = process_character_image(
            image_path,
            calibration=calibration,
            n_theta=n_theta,
            strip_height=strip_height,
            collapse_mode=collapse_mode,
        )
        strips.append(result["strip"])
        vectors.append(result["vector"])
        labels.append(Path(image_path).stem)

    raw_stack = np.vstack(vectors).astype(np.float32)
    stretched_stack = cv2.resize(
        raw_stack,
        (n_theta, int(target_height)),
        interpolation=cv2.INTER_CUBIC,
    ).astype(np.float32)
    enhanced_stack = enhance_surface(stretched_stack)

    save_image(Path(output_dir) / "raw_stack.png", normalize_surface(raw_stack))
    save_image(Path(output_dir) / "stretched_stack.png", normalize_surface(stretched_stack))
    save_image(Path(output_dir) / "enhanced_stack.png", enhanced_stack)
    save_batch_montage(Path(output_dir) / "strip_montage.png", strips, labels, cols=4)
    np.save(Path(output_dir) / "raw_vectors.npy", raw_stack)
    np.save(Path(output_dir) / "stretched_vectors.npy", stretched_stack)

    return CharacterStitchResult(
        image_paths=image_paths,
        vectors=vectors,
        raw_stack=raw_stack,
        stretched_stack=stretched_stack,
        enhanced_stack=enhanced_stack,
    )
