SIRDS Stereogram Extractor

stereodepthmap.py

"""
Stereogram Depth Map Extractor
================================
Extracts a depth map from a SIRDS (Single Image Random Dot Stereogram)
or textured stereogram by exploiting the repeating pattern.

How it works:
- Stereograms encode depth by shifting a repeating texture horizontally.
- The base "period" is the tile width. Objects closer to the viewer have
  a smaller apparent period (the repeating pattern is compressed).
- We find the base period via autocorrelation, then compute per-pixel
  disparity using a sliding cross-correlation window.
- Disparity → depth: smaller disparity = object is closer (in front).

Public domain as far as I know, written by Claude

"""

import numpy as np
from PIL import Image, ImageFilter
import argparse
import sys
from pathlib import Path


def to_grayscale(img: Image.Image) -> np.ndarray:
    """Convert PIL image to float32 grayscale array [0,1]."""
    return np.array(img.convert("L"), dtype=np.float32) / 255.0


def find_period(gray: np.ndarray, search_range=(50, 250)) -> int:
    """
    Estimate the base repeating period of the stereogram via row autocorrelation.
    We average the autocorrelation over many rows to get a robust estimate.
    The first strong peak after lag=0 is the tile period.
    """
    h, w = gray.shape
    # Use middle 60% of rows to avoid borders / watermark
    rows = gray[int(h * 0.2):int(h * 0.8), :]
    
    # Compute mean autocorrelation across rows
    mean_acf = np.zeros(w)
    for row in rows[::4]:  # sample every 4 rows for speed
        row_norm = row - row.mean()
        acf = np.correlate(row_norm, row_norm, mode='full')
        acf = acf[len(acf)//2:]  # keep positive lags
        acf /= acf[0] + 1e-9    # normalize
        mean_acf += acf[:w]
    
    mean_acf /= mean_acf[0]
    
    # Find first prominent peak in the search range
    lo, hi = search_range
    hi = min(hi, w // 3)
    region = mean_acf[lo:hi]
    
    # Find peaks: local maxima
    from scipy.signal import find_peaks
    peaks, props = find_peaks(region, height=0.3, distance=20)
    
    if len(peaks) == 0:
        # fallback: argmax
        period = int(np.argmax(region)) + lo
    else:
        period = int(peaks[0]) + lo
    
    return period


def compute_disparity_map(gray: np.ndarray, period: int,
                           window_h: int = 9, window_w: int = 9,
                           max_shift: int = 30) -> np.ndarray:
    """
    For each pixel, find the horizontal shift (disparity) that best aligns
    a window centered here with the window shifted one period to the right.
    
    In a stereogram:
      - If there's no depth object, the shift equals exactly `period`.
      - If depth is encoded, the apparent shift is period ± disparity.
    
    We search around `period` for the best cross-correlation peak.
    Returns a disparity map (float32), same shape as input.
    """
    h, w = gray.shape
    disparity = np.zeros((h, w), dtype=np.float32)
    
    ph = window_h // 2
    pw = window_w // 2
    
    # Pad image horizontally for shifting
    pad = max_shift + pw + period
    padded = np.pad(gray, ((ph, ph), (pad, pad)), mode='edge')
    
    offset = pad  # column offset in padded array
    
    for row in range(h):
        pr = row + ph  # row in padded image
        
        for col in range(pw, w - pw):
            pc = col + offset
            
            # Reference window around (row, col)
            ref = padded[pr-ph:pr+ph+1, pc-pw:pc+pw+1]
            ref_norm = ref - ref.mean()
            ref_std = ref_norm.std() + 1e-9
            
            best_score = -np.inf
            best_d = 0
            
            # Search over shifts around the period
            for d in range(-max_shift, max_shift + 1):
                shift = period + d
                if pc + shift - pw < 0 or pc + shift + pw + 1 > padded.shape[1]:
                    continue
                
                candidate = padded[pr-ph:pr+ph+1,
                                   pc+shift-pw:pc+shift+pw+1]
                cand_norm = candidate - candidate.mean()
                cand_std = cand_norm.std() + 1e-9
                
                ncc = np.sum(ref_norm * cand_norm) / (ref_std * cand_std * ref_norm.size)
                
                if ncc > best_score:
                    best_score = ncc
                    best_d = d
            
            disparity[row, col] = best_d
    
    return disparity


def compute_disparity_fast(gray: np.ndarray, period: int,
                            strip_width: int = 8,
                            max_shift: int = 35) -> np.ndarray:
    """
    Faster column-strip based disparity computation.
    Instead of per-pixel windows, we compare vertical strips of width
    `strip_width` separated by `period + d` columns.
    
    Returns disparity map (float32).
    """
    h, w = gray.shape
    disparity = np.full((h, w), np.nan, dtype=np.float32)
    
    sw = strip_width
    
    for col in range(0, w - period - max_shift - sw, sw // 2):
        ref_strip = gray[:, col:col+sw]          # left strip
        
        scores = {}
        for d in range(-max_shift, max_shift + 1):
            right_col = col + period + d
            if right_col < 0 or right_col + sw > w:
                continue
            cmp_strip = gray[:, right_col:right_col+sw]
            
            # Row-wise NCC
            diff = ref_strip - cmp_strip
            scores[d] = -np.mean(diff**2, axis=1)  # higher = better match
        
        if not scores:
            continue
        
        # Stack scores: shape (n_shifts, h)
        ds = sorted(scores.keys())
        score_matrix = np.stack([scores[d] for d in ds], axis=0)  # (n_shifts, h)
        
        # For each row, pick best shift
        best_idx = np.argmax(score_matrix, axis=0)  # (h,)
        best_d = np.array([ds[i] for i in best_idx], dtype=np.float32)
        
        # Assign to columns covered by this strip
        for c in range(col, min(col + sw, w)):
            if c < w:
                disparity[:, c] = best_d
    
    # Fill any NaN edges
    # Forward-fill columns
    for c in range(1, w):
        mask = np.isnan(disparity[:, c])
        disparity[mask, c] = disparity[mask, c-1] if c > 0 else 0
    disparity = np.nan_to_num(disparity, nan=0.0)
    
    return disparity


def smooth_disparity(disparity: np.ndarray, sigma: int = 12) -> np.ndarray:
    """Apply Gaussian smoothing to reduce noise in disparity map."""
    from scipy.ndimage import gaussian_filter
    return gaussian_filter(disparity, sigma=sigma)


def disparity_to_depthmap(disparity: np.ndarray) -> np.ndarray:
    """
    Convert disparity to a normalized 0–255 depth image.
    Convention: brighter = closer to viewer (positive disparity = object pulled forward).
    We invert so white = close, black = far (standard depth map convention).
    """
    d = disparity.copy()
    
    # Robust normalization: clip outliers
    lo, hi = np.percentile(d, 2), np.percentile(d, 98)
    if hi == lo:
        hi = lo + 1
    d = np.clip(d, lo, hi)
    d = (d - lo) / (hi - lo)
    
    # Invert: larger positive disparity = object is closer
    # (In SIRDS, objects in front create smaller period → negative d)
    # We flip so foreground appears bright
    d = 1.0 - d
    
    return (d * 255).astype(np.uint8)


def colorize_depthmap(depth_gray: np.ndarray) -> Image.Image:
    """Apply a perceptually-pleasing colormap (plasma-like) to the depth map."""
    try:
        import matplotlib.cm as cm
        colormap = cm.get_cmap('plasma')
        colored = colormap(depth_gray / 255.0)
        return Image.fromarray((colored[:, :, :3] * 255).astype(np.uint8))
    except ImportError:
        return Image.fromarray(depth_gray)


def process_stereogram(input_path: str, output_path: str = None,
                       period: int = None, strip_width: int = 8,
                       max_shift: int = 35, smooth_sigma: float = 15.0,
                       colorize: bool = True, verbose: bool = True):
    """
    Full pipeline: load → find period → compute disparity → depth map → save.
    """
    inp = Path(input_path)
    if output_path is None:
        output_path = str(inp.parent / (inp.stem + "_depthmap.png"))
    
    if verbose:
        print(f"Loading: {input_path}")
    img = Image.open(input_path)
    
    if verbose:
        print(f"  Size: {img.size}, Mode: {img.mode}")
    
    gray = to_grayscale(img)
    
    # Step 1: find period
    if period is None:
        if verbose:
            print("  Detecting tile period via autocorrelation...")
        period = find_period(gray)
        if verbose:
            print(f"  Detected period: {period}px")
    else:
        if verbose:
            print(f"  Using provided period: {period}px")
    
    # Step 2: compute disparity
    if verbose:
        print(f"  Computing disparity (strip_width={strip_width}, max_shift=±{max_shift})...")
    
    disparity = compute_disparity_fast(gray, period,
                                        strip_width=strip_width,
                                        max_shift=max_shift)
    
    # Step 3: smooth
    if verbose:
        print(f"  Smoothing (sigma={smooth_sigma})...")
    disparity = smooth_disparity(disparity, sigma=smooth_sigma)
    
    # Step 4: convert to depth image
    depth_gray = disparity_to_depthmap(disparity)
    
    # Step 5: save outputs
    # Grayscale depth map
    gray_out = str(Path(output_path).parent / (Path(output_path).stem + "_gray.png"))
    Image.fromarray(depth_gray).save(gray_out)
    if verbose:
        print(f"  Saved grayscale depth map: {gray_out}")
    
    # Colorized
    if colorize:
        color_img = colorize_depthmap(depth_gray)
        color_out = str(Path(output_path).parent / (Path(output_path).stem + "_color.png"))
        color_img.save(color_out)
        if verbose:
            print(f"  Saved colorized depth map: {color_out}")
    
    # Side-by-side comparison
    w, h = img.size
    side_by_side = Image.new("RGB", (w * 2, h))
    side_by_side.paste(img.convert("RGB"), (0, 0))
    if colorize:
        side_by_side.paste(color_img.resize((w, h)), (w, 0))
    else:
        side_by_side.paste(Image.fromarray(depth_gray).convert("RGB").resize((w, h)), (w, 0))
    side_by_side.save(output_path)
    if verbose:
        print(f"  Saved side-by-side: {output_path}")
    
    return disparity, depth_gray


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Extract a depth map from a stereogram image."
    )
    parser.add_argument("input", help="Path to stereogram image")
    parser.add_argument("-o", "--output", help="Output path for side-by-side image")
    parser.add_argument("--period", type=int, default=None,
                        help="Override auto-detected tile period (pixels)")
    parser.add_argument("--strip-width", type=int, default=8,
                        help="Comparison strip width (default: 8)")
    parser.add_argument("--max-shift", type=int, default=35,
                        help="Max disparity to search ±N pixels (default: 35)")
    parser.add_argument("--smooth", type=float, default=15.0,
                        help="Gaussian smoothing sigma (default: 15)")
    parser.add_argument("--no-color", action="store_true",
                        help="Skip colorized output")
    args = parser.parse_args()
    
    process_stereogram(
        args.input,
        output_path=args.output,
        period=args.period,
        strip_width=args.strip_width,
        max_shift=args.max_shift,
        smooth_sigma=args.smooth,
        colorize=not args.no_color,
    )