2d_localize.py

import argparse
import os
import yaml
import numpy as np
import torch
from PIL import Image

from nn.marker_detection import MarkerDetectionInference
from nn.marker_recognition import (
    MarkerRecognitionInference,
    gen_garment_labels,
    load_marker_info,
)
from nn.uv_localization import UVLocalizationInference
from nn.recover_mesh_2d import RecoverMesh2DPipeline


_STAGE_DIRS = {
    1: "stage1_detection",
    2: "stage2_recognition",
    3: "stage3_uv",
    4: "stage4_mesh2d",
}


def _stage_dir(output_root, stage):
    return os.path.join(output_root, _STAGE_DIRS[stage])


def _setup_dirs(output_root, garment_names, vis_det, vis_recog, vis_uv):
    if vis_det:
        os.makedirs(os.path.join(_stage_dir(output_root, 1), "vis"), exist_ok=True)
    if vis_recog:
        os.makedirs(os.path.join(_stage_dir(output_root, 2), "vis"), exist_ok=True)
    if vis_uv:
        os.makedirs(os.path.join(_stage_dir(output_root, 3), "vis"), exist_ok=True)
        os.makedirs(os.path.join(_stage_dir(output_root, 4), "vis"), exist_ok=True)
    coords_2d = os.path.join(output_root, "coords_2d")
    for g in garment_names:
        os.makedirs(os.path.join(coords_2d, g), exist_ok=True)
    return coords_2d


def _vis1(detection_results, stage1_dir):
    import cv2
    from visualize import visualize_markers
    vis_dir = os.path.join(stage1_dir, "vis")
    for camera, frame, markers, img in detection_results:
        frame = int(frame)
        img_np = img.cpu().numpy() if isinstance(img, torch.Tensor) else img
        img_bgr = cv2.cvtColor(
            (img_np.transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR
        )
        visualize_markers(img_bgr, markers)
        cv2.imwrite(os.path.join(vis_dir, f"frame_{frame:05d}_cam{camera}.png"), img_bgr)


def _vis2(detection_results, garment_labels, garment_names, stage2_dir):
    import cv2
    from visualize import visualize_recognition
    vis_dir = os.path.join(stage2_dir, "vis")
    for i, (camera, frame, _, img) in enumerate(detection_results):
        frame = int(frame)
        img_np = img.cpu().numpy() if isinstance(img, torch.Tensor) else img
        img_bgr = cv2.cvtColor(
            (img_np.transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR
        )
        for g in range(len(garment_names)):
            _, _, g_markers, g_labels, _ = garment_labels[g][i]
            if len(g_markers) > 0:
                visualize_recognition(img_bgr, g_markers, g_labels)
        cv2.imwrite(os.path.join(vis_dir, f"frame_{frame:05d}_cam{camera}.png"), img_bgr)


def _vis3(all_uv_results, garment_names, stage3_dir):
    from visualize import visualize_uv
    vis_dir = os.path.join(stage3_dir, "vis")
    for g, gname in enumerate(garment_names):
        for uv_result in all_uv_results[g]:
            camera, frame, outputs, warped_imgs, crop_ids, _ = uv_result
            frame = int(frame)
            stem = f"frame_{frame:05d}_{gname}_cam{camera}"
            outputs_np = outputs.cpu().numpy() if isinstance(outputs, torch.Tensor) else outputs
            crops = warped_imgs.cpu() if isinstance(warped_imgs, torch.Tensor) else torch.from_numpy(warped_imgs)
            visualize_uv(crops, outputs_np, crop_ids, stem, vis_dir)


def _vis4(recover_result, image_dir, stage4_dir):
    import cv2
    from visualize import visualize_recovered_2d
    vis_dir = os.path.join(stage4_dir, "vis")
    camera, frame, recovered_coord2D, _, _ = recover_result
    frame = int(frame)
    img_path = os.path.join(image_dir, camera, f"{camera}{frame:05d}.png")
    img = np.array(Image.open(img_path).convert("RGB"))
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    visualize_recovered_2d(img_bgr, recovered_coord2D)
    cv2.imwrite(os.path.join(vis_dir, f"frame_{frame:05d}_cam{camera}.png"), img_bgr)


def run_2d_localize(config, vis_det=False, vis_recog=False, vis_uv=False):
    pipeline_cfg = config["pipeline"]
    image_dir = pipeline_cfg["image_dir"]
    cameras = pipeline_cfg["cameras"]
    start_frame = pipeline_cfg["start_frame"]
    end_frame = pipeline_cfg["end_frame"]
    cuda = pipeline_cfg["cuda"]
    output_root = pipeline_cfg["output_root"]
    recog_threshold = pipeline_cfg.get("recog_threshold", 0.99999)

    garment_names = pipeline_cfg["garments"]
    garment_configs = [config["garments"][g] for g in garment_names]

    coords_2d_dir = _setup_dirs(output_root, garment_names, vis_det, vis_recog, vis_uv)
    s1_dir = _stage_dir(output_root, 1)
    s2_dir = _stage_dir(output_root, 2)
    s3_dir = _stage_dir(output_root, 3)
    s4_dir = _stage_dir(output_root, 4)

    print("Loading detection model...")
    detection = MarkerDetectionInference(
        model_path=pipeline_cfg["detection_model_path"], cuda=cuda
    )

    print("Loading recognition models...")
    all_recognition_models = []
    all_marker_mappings = []
    for gcfg in garment_configs:
        recognition = MarkerRecognitionInference(
            model_path=gcfg["recognition_model_path"],
            cuda=cuda,
            inner_scale=gcfg["inner_scale"],
            outer_scale=gcfg["outer_scale"],
            markerW=gcfg.get("markerW", 24),
            markerH=gcfg.get("markerH", 24),
        )
        all_recognition_models.append(recognition)
        _, marker_mappings = load_marker_info(
            gcfg["marker_info_path"], recognition.n_markers
        )
        all_marker_mappings.append(marker_mappings)

    print("Loading UV localization models...")
    all_uv_models = []
    for gcfg in garment_configs:
        uv = UVLocalizationInference(
            model_path=gcfg["uv_model_path"],
            crop_info_path=gcfg["crop_info_path"],
            cuda=cuda,
        )
        all_uv_models.append(uv)

    print("Loading mesh recovery pipelines...")
    all_mesh_recovery = []
    for gcfg in garment_configs:
        mesh_recovery = RecoverMesh2DPipeline(
            markerW=gcfg.get("markerW", 24),
            markerH=gcfg.get("markerH", 24),
            crop_info_path=gcfg["crop_info_path"],
            mesh_crop_info_path=gcfg["mesh_crop_info_path"],
            mesh_npz_path=gcfg["mesh_npz_path"],
            cuda=cuda,
        )
        all_mesh_recovery.append(mesh_recovery)

    print(f"Running 2D localization on frames {start_frame}–{end_frame - 1}...")
    for frame in range(start_frame, end_frame):
        detection.load_images(
            img_dir=image_dir,
            start_frame=frame,
            end_frame=frame + 1,
            cameras=cameras,
        )
        detection_results = detection.infer_markers(batch_size=1)
        if vis_det:
            _vis1(detection_results, s1_dir)

        all_recog_results = [
            m.recognize_marker(detection_results=detection_results, batch_size=256)
            for m in all_recognition_models
        ]
        garment_labels = gen_garment_labels(
            detection_results, all_recog_results, all_marker_mappings, threshold=recog_threshold
        )
        if vis_recog:
            _vis2(detection_results, garment_labels, garment_names, s2_dir)

        all_uv_results = [
            uv.localize_uv(
                recognition_results=garment_labels[i],
                detection_results=detection_results,
                batch_size=256,
            )
            for i, uv in enumerate(all_uv_models)
        ]
        if vis_uv:
            _vis3(all_uv_results, garment_names, s3_dir)

        for i, gcfg in enumerate(garment_configs):
            garment_out = os.path.join(coords_2d_dir, gcfg["name"])
            for uv_result in all_uv_results[i]:
                recover_result = all_mesh_recovery[i].recover_mesh_2d_warp(
                    uv_result=uv_result, threshold=0.9, std_threshold=1.0
                )
                all_mesh_recovery[i].output_coords_2d(recover_result, garment_out)
                if vis_uv:
                    _vis4(recover_result, image_dir, s4_dir)

    print("Done.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="ClothCap stages 1–4: marker detection, recognition, UV localization, 2D mesh recovery"
    )
    parser.add_argument("--config", required=True, help="Path to YAML config file")
    parser.add_argument(
        "--sequence", default=None, metavar="NAME",
        help="Sequence name from the config's sequences block"
    )
    parser.add_argument("--vis-det", action="store_true", help="Write detection visualizations")
    parser.add_argument("--vis-recog", action="store_true", help="Write recognition visualizations")
    parser.add_argument("--vis-uv", action="store_true", help="Write UV localization and 2D mesh recovery visualizations")
    parser.add_argument("--cuda", type=int, default=None,
                        help="CUDA device index override from config; use -1 for CPU")
    args = parser.parse_args()

    with open(args.config, encoding="utf-8") as f:
        config = yaml.safe_load(f)

    if args.sequence is not None:
        sequences = config.get("sequences", {})
        if args.sequence not in sequences:
            raise SystemExit(f"Unknown sequence '{args.sequence}'. Available: {list(sequences)}")
        config["pipeline"].update(sequences[args.sequence])
    if args.cuda is not None:
        config["pipeline"]["cuda"] = args.cuda

    run_2d_localize(
        config,
        vis_det=args.vis_det,
        vis_recog=args.vis_recog,
        vis_uv=args.vis_uv,
    )