blend.py

#!/usr/bin/env python3
import argparse
from pathlib import Path

import cv2
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Canny edge mix: edge regions use refine image, "
            "non-edge regions use weighted average."
        )
    )
    parser.add_argument("--base-dir", type=str, required=True, help="Directory of base images.")
    parser.add_argument("--refine-dir", type=str, required=True, help="Directory of refine images.")
    parser.add_argument("--out-dir", type=str, required=True, help="Output root directory.")
    parser.add_argument(
        "--exts",
        type=str,
        nargs="+",
        default=["png", "jpg", "jpeg", "bmp", "webp"],
        help="Image extensions to match (case-insensitive).",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="Only process first N matched pairs (0 means all).",
    )

    # Keep these five tunable parameters only.
    parser.add_argument(
        "--nonedge-alpha",
        type=float,
        default=0.5,
        help="Base-image weight for non-edge averaging (0~1).",
    )
    parser.add_argument("--canny-low", type=int, default=80, help="Canny lower threshold.")
    parser.add_argument("--canny-high", type=int, default=160, help="Canny upper threshold.")
    parser.add_argument(
        "--canny-dilate",
        type=int,
        default=1,
        help="Dilate iterations for edge mask.",
    )
    parser.add_argument(
        "--edge-feather",
        type=int,
        default=1,
        help="Gaussian feather radius for edge mask (0 means hard edge).",
    )
    return parser.parse_args()


def build_pairs(base_dir: Path, refine_dir: Path, exts):
    exts_set = {e.lower().lstrip(".") for e in exts}
    base_map = {}
    for p in base_dir.rglob("*"):
        if p.is_file() and p.suffix.lower().lstrip(".") in exts_set:
            rel = p.relative_to(base_dir).as_posix()
            base_map[rel] = p

    pairs = []
    missing = 0
    for rel, base_path in sorted(base_map.items()):
        refine_path = refine_dir / rel
        if refine_path.exists():
            pairs.append((rel, base_path, refine_path))
        else:
            missing += 1
    return pairs, missing


def canny_edge_mask(base_bgr_u8, canny_low, canny_high, dilate_iter=1, feather_radius=1):
    gray = cv2.cvtColor(base_bgr_u8, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0.0)
    edge = cv2.Canny(gray, int(canny_low), int(canny_high))
    if dilate_iter > 0:
        kernel = np.ones((3, 3), np.uint8)
        edge = cv2.dilate(edge, kernel, iterations=int(dilate_iter))
    edge01 = edge.astype(np.float32) / 255.0
    if feather_radius > 0:
        k = int(feather_radius) * 2 + 1
        edge01 = cv2.GaussianBlur(edge01, (k, k), 0.0)
    edge01 = np.clip(edge01, 0.0, 1.0)
    return edge01, edge


def blend_with_base_weight(base_bgr_u8, refine_bgr_u8, base_weight_01):
    base = base_bgr_u8.astype(np.float32) / 255.0
    refine = refine_bgr_u8.astype(np.float32) / 255.0
    w = base_weight_01[:, :, None]
    out = w * base + (1.0 - w) * refine
    out = np.clip(out * 255.0 + 0.5, 0.0, 255.0).astype(np.uint8)
    return out


def ensure_parent(path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)


def main():
    args = parse_args()
    base_dir = Path(args.base_dir)
    refine_dir = Path(args.refine_dir)
    out_dir = Path(args.out_dir)

    out_blend = out_dir / "blended"
    out_mask = out_dir / "mask01"
    out_mask_vis = out_dir / "mask_vis"
    out_blend.mkdir(parents=True, exist_ok=True)
    out_mask.mkdir(parents=True, exist_ok=True)
    out_mask_vis.mkdir(parents=True, exist_ok=True)

    pairs, missing = build_pairs(base_dir, refine_dir, args.exts)
    if args.limit > 0:
        pairs = pairs[: args.limit]
    print(f"Found pairs: {len(pairs)}, missing in refine: {missing}")
    if not pairs:
        raise RuntimeError("No matching image pairs found.")

    fail_count = 0
    alpha_nonedge = float(np.clip(args.nonedge_alpha, 0.0, 1.0))
    for idx, (rel, base_path, refine_path) in enumerate(pairs, 1):
        base = cv2.imread(str(base_path), cv2.IMREAD_COLOR)
        refine = cv2.imread(str(refine_path), cv2.IMREAD_COLOR)
        if base is None or refine is None:
            print(f"[WARN] Failed to read: {rel}")
            fail_count += 1
            continue
        if base.shape != refine.shape:
            print(f"[WARN] Shape mismatch, skip: {rel}, {base.shape} vs {refine.shape}")
            fail_count += 1
            continue

        edge01, edge_bin = canny_edge_mask(
            base,
            args.canny_low,
            args.canny_high,
            dilate_iter=args.canny_dilate,
            feather_radius=args.edge_feather,
        )

        # Edge: use refine directly (base weight = 0).
        # Non-edge: weighted average by nonedge-alpha.
        base_weight = np.clip((1.0 - edge01) * alpha_nonedge, 0.0, 1.0)
        out = blend_with_base_weight(base, refine, base_weight)

        rel_path = Path(rel)
        p_blend = out_blend / rel_path
        p_mask = out_mask / rel_path
        p_mvis = out_mask_vis / rel_path
        ensure_parent(p_blend)
        ensure_parent(p_mask)
        ensure_parent(p_mvis)

        cv2.imwrite(str(p_blend), out)
        cv2.imwrite(str(p_mask), (base_weight * 255.0 + 0.5).astype(np.uint8))

        heat = cv2.applyColorMap(edge_bin.astype(np.uint8), cv2.COLORMAP_JET)
        vis = cv2.addWeighted(refine, 0.78, heat, 0.22, 0.0)
        cv2.imwrite(str(p_mvis), vis)

        if idx % 100 == 0 or idx == len(pairs):
            print(f"Processed {idx}/{len(pairs)}")

    print(f"Done. failures={fail_count}, output={out_dir}")


if __name__ == "__main__":
    main()