ioandanielc/sph_dataset

1#!/usr/bin/env python32"""Apply border recoloring to side/front frames; copy top frames as-is.3 4Input:  final_data/sim_NNNNN/frames/{front,side,top}/frame_NNNNN.png5Output: final_data_processed/sim_NNNNN/frames/{front,side,top}/frame_NNNNN.png6"""7from __future__ import annotations8 9import os10import shutil11from pathlib import Path12 13import numpy as np14from PIL import Image15 16# ── constants (from prepare.py / cvae_sph2img) ───────────────────────────────17PURE_GREEN           = np.array([0, 197, 0],   dtype=np.uint8)18PURE_BLUE            = np.array([0, 0, 189],   dtype=np.uint8)19PURE_BLACK           = np.array([0, 0, 0],     dtype=np.uint8)20PURE_GRAY            = np.array([98, 93, 90],  dtype=np.uint8)21BORDER_THICKNESS_PX  = 1022SIDE_BOTTOM_HEIGHT_PX = 18823SIDE_BUFFER          = 324 25# ── image processing ──────────────────────────────────────────────────────────26 27def detect_side_columns(rgb: np.ndarray) -> tuple[int, int]:28    h, w, _ = rgb.shape29    if h == 0 or w == 0:30        return (0, 0)31    probe_y = max(0, h - SIDE_BOTTOM_HEIGHT_PX)32    row = rgb[probe_y, :, :3]33    mask = np.all(row == PURE_BLACK, axis=1) | np.all(row == PURE_GRAY, axis=1)34    left = 035    while left < w and mask[left]:36        left += 137    right = 038    idx = w - 139    while idx >= 0 and mask[idx]:40        right += 141        idx -= 142    return (left, right)43 44 45def fixed_border_recolor(rgb: np.ndarray, left_columns: int, right_columns: int) -> np.ndarray:46    out = rgb[:, :, :3].copy()47    h, w, _ = out.shape48    bt  = min(BORDER_THICKNESS_PX, h, w)49    sbh = min(SIDE_BOTTOM_HEIGHT_PX, h)50    out[:bt, :, :]  = PURE_BLUE51    out[h - bt:, :, :] = PURE_GREEN52    if left_columns > 0:53        lw = min(left_columns, w)54        out[:, :lw, :]          = PURE_BLUE55        out[h - sbh:, :lw, :]  = PURE_GREEN56    if right_columns > 0:57        rw = min(right_columns, w)58        out[:, w - rw:, :]         = PURE_BLUE59        out[h - sbh:, w - rw:, :]  = PURE_GREEN60    return out61 62 63def process_image(src: Path, dst: Path, is_side: bool) -> None:64    rgb = np.array(Image.open(src).convert("RGB"))65    h, w, _ = rgb.shape66    default_cols = min(BORDER_THICKNESS_PX, h, w)67    if is_side:68        left, right = detect_side_columns(rgb)69        if left  > 0: left  = min(w, left  + SIDE_BUFFER)70        if right > 0: right = min(w, right + SIDE_BUFFER)71    else:72        left, right = default_cols, default_cols73    out = fixed_border_recolor(rgb, left, right)74    Image.fromarray(out).save(dst)75 76 77# ── main ──────────────────────────────────────────────────────────────────────78 79def main() -> None:80    here        = Path(__file__).parent81    src_root    = here / "final_data"82    dst_root    = here / "final_data_processed"83 84    sim_dirs = sorted(src_root.iterdir())85    print(f"Found {len(sim_dirs)} simulations")86 87    for sim_dir in sim_dirs:88        if not sim_dir.is_dir():89            continue90        frames_src = sim_dir / "frames"91        if not frames_src.is_dir():92            print(f"  SKIP {sim_dir.name} — no frames/ directory")93            continue94 95        for view in ("front", "side", "top"):96            view_src = frames_src / view97            view_dst = dst_root / sim_dir.name / "frames" / view98            if not view_src.is_dir():99                continue100            view_dst.mkdir(parents=True, exist_ok=True)101 102            for png in sorted(view_src.glob("*.png")):103                dst_path = view_dst / png.name104                if view == "top":105                    shutil.copy2(png, dst_path)106                else:107                    process_image(png, dst_path, is_side=(view == "side"))108 109        print(f"  {sim_dir.name} done")110 111    print("All done.")112 113 114if __name__ == "__main__":115    main()116