SignerX/SignVerse-2M

1#!/usr/bin/env python32 3import argparse4import importlib.util5import math6import shutil7import subprocess8import tempfile9from pathlib import Path10from typing import Dict, Iterable, List11 12import cv213import matplotlib14import numpy as np15from PIL import Image16 17 18REPO_ROOT = Path(__file__).resolve().parents[1]19 20import sys21 22if str(REPO_ROOT) not in sys.path:23    sys.path.insert(0, str(REPO_ROOT))24 25from utils.draw_dw_lib import draw_pose26 27 28VIDEO_EXTENSIONS = (".mp4", ".mkv", ".mov", ".webm")29EPS = 0.0130STABLE_SIGNER_OPENPOSE_PATH = Path(31    "/research/cbim/vast/sf895/code/SignerX-inference-webui/plugins/StableSigner/easy_dwpose/draw/openpose.py"32)33_STABLE_SIGNER_OPENPOSE_DRAW = None34 35 36def parse_args() -> argparse.Namespace:37    parser = argparse.ArgumentParser(description="Visualize Sign-DWPose NPZ outputs.")38    parser.add_argument("--video-dir", type=Path, required=True, help="Dataset video directory, e.g. dataset/<video_id>")39    parser.add_argument("--npz-dir", type=Path, default=None, help="Optional NPZ directory override")40    parser.add_argument("--raw-video", type=Path, default=None, help="Optional raw video path for overlay rendering")41    parser.add_argument("--fps", type=int, default=24, help="Visualization FPS")42    parser.add_argument("--max-frames", type=int, default=None, help="Limit the number of frames to render")43    parser.add_argument(44        "--draw-style",45        choices=("controlnext", "openpose", "dwpose"),46        default="controlnext",47        help="Rendering style. dwpose is kept as an alias of controlnext.",48    )49    parser.add_argument("--conf-threshold", type=float, default=0.6, help="Confidence threshold for openpose filtering")50    parser.add_argument(51        "--frame-indices",52        default="1,2,3,4",53        help="Comma-separated 1-based frame indices for standalone single-frame previews",54    )55    parser.add_argument(56        "--output-dir",57        type=Path,58        default=None,59        help="Visualization output directory. Defaults to <video-dir>/visualization_dwpose",60    )61    parser.add_argument("--force", action="store_true", help="Overwrite existing visualization outputs")62    return parser.parse_args()63 64 65def parse_frame_indices(value: str) -> List[int]:66    indices: List[int] = []67    for item in value.split(","):68        item = item.strip()69        if not item:70            continue71        index = int(item)72        if index > 0:73            indices.append(index)74    return sorted(set(indices))75 76 77def normalize_draw_style(value: str) -> str:78    return "controlnext" if value == "dwpose" else value79 80 81def get_stablesigner_openpose_draw():82    global _STABLE_SIGNER_OPENPOSE_DRAW  # noqa: PLW060383    if _STABLE_SIGNER_OPENPOSE_DRAW is not None:84        return _STABLE_SIGNER_OPENPOSE_DRAW85    if not STABLE_SIGNER_OPENPOSE_PATH.exists():86        return None87    spec = importlib.util.spec_from_file_location("stablesigner_openpose_draw", STABLE_SIGNER_OPENPOSE_PATH)88    if spec is None or spec.loader is None:89        return None90    module = importlib.util.module_from_spec(spec)91    spec.loader.exec_module(module)92    _STABLE_SIGNER_OPENPOSE_DRAW = getattr(module, "draw_pose", None)93    return _STABLE_SIGNER_OPENPOSE_DRAW94 95 96def load_npz_frame(npz_path: Path, aggregated_index: int = 0) -> Dict[str, object]:97    payload = np.load(npz_path, allow_pickle=True)98    if "frame_payloads" in payload.files:99        frame_payloads = payload["frame_payloads"]100        if aggregated_index >= len(frame_payloads):101            raise IndexError(f"Aggregated frame index {aggregated_index} out of range for {npz_path}")102        payload_dict = frame_payloads[aggregated_index]103        if hasattr(payload_dict, "item"):104            payload_dict = payload_dict.item()105        frame: Dict[str, object] = {}106        frame["num_persons"] = int(payload_dict["num_persons"])107        frame["frame_width"] = int(payload_dict["frame_width"])108        frame["frame_height"] = int(payload_dict["frame_height"])109        source = payload_dict110    else:111        frame = {}112        frame["num_persons"] = int(payload["num_persons"])113        frame["frame_width"] = int(payload["frame_width"])114        frame["frame_height"] = int(payload["frame_height"])115        source = payload116 117    for person_idx in range(frame["num_persons"]):118        source_prefix = f"person_{person_idx:03d}"119        target_prefix = f"person_{person_idx}"120        person_data: Dict[str, np.ndarray] = {}121        for suffix in (122            "body_keypoints",123            "body_scores",124            "face_keypoints",125            "face_scores",126            "left_hand_keypoints",127            "left_hand_scores",128            "right_hand_keypoints",129            "right_hand_scores",130        ):131            key = f"{source_prefix}_{suffix}"132            if key in source:133                person_data[suffix] = source[key]134        if person_data:135            frame[target_prefix] = person_data136    return frame137 138 139def to_openpose_frame(frame: Dict[str, object]) -> Dict[str, np.ndarray]:140    num_persons = int(frame["num_persons"])141    bodies: List[np.ndarray] = []142    body_scores: List[np.ndarray] = []143    hands: List[np.ndarray] = []144    hand_scores: List[np.ndarray] = []145    faces: List[np.ndarray] = []146    face_scores: List[np.ndarray] = []147 148    for person_idx in range(num_persons):149        person = frame.get(f"person_{person_idx}")150        if not isinstance(person, dict):151            continue152        bodies.append(np.asarray(person["body_keypoints"], dtype=np.float32))153        body_scores.append(np.asarray(person["body_scores"], dtype=np.float32))154        hands.extend(155            [156                np.asarray(person["left_hand_keypoints"], dtype=np.float32),157                np.asarray(person["right_hand_keypoints"], dtype=np.float32),158            ]159        )160        hand_scores.extend(161            [162                np.asarray(person["left_hand_scores"], dtype=np.float32),163                np.asarray(person["right_hand_scores"], dtype=np.float32),164            ]165        )166        faces.append(np.asarray(person["face_keypoints"], dtype=np.float32))167        face_scores.append(np.asarray(person["face_scores"], dtype=np.float32))168 169    if bodies:170        stacked_bodies = np.vstack(bodies)171        stacked_subset = np.vstack(body_scores)172    else:173        stacked_bodies = np.zeros((0, 2), dtype=np.float32)174        stacked_subset = np.zeros((0, 18), dtype=np.float32)175 176    return {177        "bodies": stacked_bodies,178        "body_scores": stacked_subset,179        "hands": np.asarray(hands, dtype=np.float32) if hands else np.zeros((0, 21, 2), dtype=np.float32),180        "hands_scores": np.asarray(hand_scores, dtype=np.float32) if hand_scores else np.zeros((0, 21), dtype=np.float32),181        "faces": np.asarray(faces, dtype=np.float32) if faces else np.zeros((0, 68, 2), dtype=np.float32),182        "faces_scores": np.asarray(face_scores, dtype=np.float32) if face_scores else np.zeros((0, 68), dtype=np.float32),183    }184 185 186def filter_pose_for_openpose(frame: Dict[str, np.ndarray], conf_threshold: float, update_subset: bool) -> Dict[str, np.ndarray]:187    filtered = {key: np.array(value, copy=True) for key, value in frame.items()}188 189    bodies = filtered.get("bodies", None)190    body_scores = filtered.get("body_scores", None)191    if bodies is not None:192        bodies = bodies.copy()193        min_valid = 1e-6194        coord_mask = (bodies[:, 0] > min_valid) & (bodies[:, 1] > min_valid)195 196        conf_mask = None197        if body_scores is not None:198            scores = np.array(body_scores, copy=False)199            score_vec = scores.reshape(-1) if scores.ndim == 2 else scores200            score_vec = score_vec.astype(float)201            conf_mask = score_vec < conf_threshold202            if conf_mask.shape[0] < bodies.shape[0]:203                conf_mask = np.pad(conf_mask, (0, bodies.shape[0] - conf_mask.shape[0]), constant_values=False)204            elif conf_mask.shape[0] > bodies.shape[0]:205                conf_mask = conf_mask[: bodies.shape[0]]206        valid_mask = coord_mask if conf_mask is None else (coord_mask & (~conf_mask))207        bodies[~valid_mask, :] = 0208        filtered["bodies"] = bodies209 210        if update_subset:211            if body_scores is not None:212                subset = np.array(body_scores, copy=True)213                if subset.ndim == 1:214                    subset = subset.reshape(1, -1)215            else:216                subset = np.arange(bodies.shape[0], dtype=float).reshape(1, -1)217            if subset.shape[1] < bodies.shape[0]:218                subset = np.pad(subset, ((0, 0), (0, bodies.shape[0] - subset.shape[1])), constant_values=-1)219            elif subset.shape[1] > bodies.shape[0]:220                subset = subset[:, : bodies.shape[0]]221            subset[:, ~valid_mask] = -1222            filtered["body_scores"] = subset223 224    hands = filtered.get("hands", None)225    hand_scores = filtered.get("hands_scores", None)226    if hands is not None and hand_scores is not None:227        scores = np.array(hand_scores)228        hands = hands.copy()229        if hands.ndim == 3 and scores.ndim == 2:230            for hand_index in range(hands.shape[0]):231                mask = (scores[hand_index] < conf_threshold) | (scores[hand_index] <= 0)232                hands[hand_index][mask, :] = 0233        filtered["hands"] = hands234 235    faces = filtered.get("faces", None)236    face_scores = filtered.get("faces_scores", None)237    if faces is not None and face_scores is not None:238        scores = np.array(face_scores)239        faces = faces.copy()240        if faces.ndim == 3 and scores.ndim == 2:241            for face_index in range(faces.shape[0]):242                mask = (scores[face_index] < conf_threshold) | (scores[face_index] <= 0)243                faces[face_index][mask, :] = 0244        filtered["faces"] = faces245 246    return filtered247 248 249def draw_openpose_body(canvas: np.ndarray, candidate: np.ndarray, subset: np.ndarray, score: np.ndarray, conf_threshold: float) -> np.ndarray:250    height, width, _ = canvas.shape251    limb_seq = [252        [2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11],253        [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18],254    ]255    colors = [256        [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],257        [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],258        [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85],259    ]260 261    for limb_index in range(17):262        for person_index in range(len(subset)):263            index = subset[person_index][np.array(limb_seq[limb_index]) - 1]264            if -1 in index:265                continue266            confidence = score[person_index][np.array(limb_seq[limb_index]) - 1]267            if confidence[0] < conf_threshold or confidence[1] < conf_threshold:268                continue269            coords = candidate[index.astype(int)]270            if np.any(coords <= EPS):271                continue272            y_coords = coords[:, 0] * float(width)273            x_coords = coords[:, 1] * float(height)274            mean_x = np.mean(x_coords)275            mean_y = np.mean(y_coords)276            length = ((x_coords[0] - x_coords[1]) ** 2 + (y_coords[0] - y_coords[1]) ** 2) ** 0.5277            angle = math.degrees(math.atan2(x_coords[0] - x_coords[1], y_coords[0] - y_coords[1]))278            polygon = cv2.ellipse2Poly((int(mean_y), int(mean_x)), (int(length / 2), 4), int(angle), 0, 360, 1)279            cv2.fillConvexPoly(canvas, polygon, colors[limb_index])280 281    canvas = (canvas * 0.6).astype(np.uint8)282    for keypoint_index in range(18):283        for person_index in range(len(subset)):284            index = int(subset[person_index][keypoint_index])285            if index == -1 or score[person_index][keypoint_index] < conf_threshold:286                continue287            x_value, y_value = candidate[index][0:2]288            cv2.circle(canvas, (int(x_value * width), int(y_value * height)), 4, colors[keypoint_index], thickness=-1)289    return canvas290 291 292def draw_openpose_hands(canvas: np.ndarray, hand_peaks: np.ndarray, hand_scores: np.ndarray, conf_threshold: float) -> np.ndarray:293    height, width, _ = canvas.shape294    edges = [295        [0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10],296        [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20],297    ]298    for hand_index, peaks in enumerate(hand_peaks):299        scores = hand_scores[hand_index] if len(hand_scores) > hand_index else None300        for edge_index, edge in enumerate(edges):301            x1, y1 = peaks[edge[0]]302            x2, y2 = peaks[edge[1]]303            if scores is not None and (scores[edge[0]] < conf_threshold or scores[edge[1]] < conf_threshold):304                continue305            x1 = int(x1 * width)306            y1 = int(y1 * height)307            x2 = int(x2 * width)308            y2 = int(y2 * height)309            if x1 > EPS and y1 > EPS and x2 > EPS and y2 > EPS:310                cv2.line(311                    canvas,312                    (x1, y1),313                    (x2, y2),314                    matplotlib.colors.hsv_to_rgb([edge_index / float(len(edges)), 1.0, 1.0]) * 255,315                    thickness=2,316                )317        for point_index, point in enumerate(peaks):318            if scores is not None and scores[point_index] < conf_threshold:319                continue320            x_value = int(point[0] * width)321            y_value = int(point[1] * height)322            if x_value > EPS and y_value > EPS:323                cv2.circle(canvas, (x_value, y_value), 4, (0, 0, 255), thickness=-1)324    return canvas325 326 327def draw_openpose_faces(canvas: np.ndarray, face_points: np.ndarray, face_scores: np.ndarray, conf_threshold: float) -> np.ndarray:328    height, width, _ = canvas.shape329    for face_index, points in enumerate(face_points):330        scores = face_scores[face_index] if len(face_scores) > face_index else None331        for point_index, point in enumerate(points):332            if scores is not None and scores[point_index] < conf_threshold:333                continue334            x_value = int(point[0] * width)335            y_value = int(point[1] * height)336            if x_value > EPS and y_value > EPS:337                cv2.circle(canvas, (x_value, y_value), 3, (255, 255, 255), thickness=-1)338    return canvas339 340 341def draw_openpose_frame(frame: Dict[str, np.ndarray], width: int, height: int, conf_threshold: float) -> Image.Image:342    draw_func = get_stablesigner_openpose_draw()343    if draw_func is not None:344        canvas = draw_func(345            pose=frame,346            height=height,347            width=width,348            include_face=True,349            include_hands=True,350            conf_threshold=conf_threshold,351        )352        return Image.fromarray(canvas, "RGB")353 354    canvas = np.zeros((height, width, 3), dtype=np.uint8)355    bodies = frame["bodies"]356    subset = frame.get("body_scores", np.zeros((1, 18), dtype=np.float32))357    if subset.ndim == 1:358        subset = subset.reshape(1, -1)359    canvas = draw_openpose_body(canvas, bodies, subset, subset, conf_threshold)360    if len(frame.get("faces", [])) > 0:361        canvas = draw_openpose_faces(canvas, frame["faces"], frame.get("faces_scores", np.zeros((0, 68))), conf_threshold)362    if len(frame.get("hands", [])) > 0:363        canvas = draw_openpose_hands(canvas, frame["hands"], frame.get("hands_scores", np.zeros((0, 21))), conf_threshold)364    return Image.fromarray(canvas, "RGB")365 366 367def render_pose_image(frame: Dict[str, object], draw_style: str, transparent: bool, conf_threshold: float) -> Image.Image:368    width = int(frame["frame_width"])369    height = int(frame["frame_height"])370    if draw_style == "openpose":371        openpose_frame = filter_pose_for_openpose(372            to_openpose_frame(frame),373            conf_threshold=conf_threshold,374            update_subset=True,375        )376        image = draw_openpose_frame(openpose_frame, width, height, conf_threshold)377        if not transparent:378            return image379        rgba = image.convert("RGBA")380        alpha = np.where(np.array(image).sum(axis=2) > 0, 255, 0).astype(np.uint8)381        rgba.putalpha(Image.fromarray(alpha, "L"))382        return rgba383 384    rendered = draw_pose(385        frame,386        H=height,387        W=width,388        include_body=True,389        include_hand=True,390        include_face=True,391        transparent=transparent,392    )393    rendered = np.transpose(rendered, (1, 2, 0))394    if rendered.dtype != np.uint8:395        rendered = np.clip(rendered * 255.0, 0, 255).astype(np.uint8)396    return Image.fromarray(rendered, "RGBA" if transparent else "RGB")397 398 399def save_frame_previews(npz_paths: Iterable[Path], single_frame_dir: Path, draw_style: str, conf_threshold: float) -> None:400    single_frame_dir.mkdir(parents=True, exist_ok=True)401    for preview_index, npz_path in enumerate(npz_paths, start=1):402        frame = load_npz_frame(npz_path, aggregated_index=preview_index - 1 if npz_path.name == "poses.npz" else 0)403        image = render_pose_image(frame, draw_style=draw_style, transparent=False, conf_threshold=conf_threshold)404        image.save(single_frame_dir / f"{npz_path.stem}.png")405 406 407def render_pose_frames(npz_paths: List[Path], pose_frame_dir: Path, draw_style: str, conf_threshold: float) -> None:408    pose_frame_dir.mkdir(parents=True, exist_ok=True)409    total = len(npz_paths)410    for index, npz_path in enumerate(npz_paths, start=1):411        frame = load_npz_frame(npz_path, aggregated_index=index - 1 if npz_path.name == "poses.npz" else 0)412        image = render_pose_image(frame, draw_style=draw_style, transparent=False, conf_threshold=conf_threshold)413        image.save(pose_frame_dir / f"{npz_path.stem}.png")414        if index == 1 or index % 100 == 0 or index == total:415            print(f"Rendered pose frame {index}/{total}: {npz_path.name}")416 417 418def create_video_from_frames(frame_dir: Path, output_path: Path, fps: int) -> None:419    if not any(frame_dir.glob("*.png")):420        return421    command = [422        "ffmpeg",423        "-hide_banner",424        "-loglevel",425        "error",426        "-y",427        "-framerate",428        str(fps),429        "-i",430        str(frame_dir / "%08d.png"),431        "-c:v",432        "libx264",433        "-pix_fmt",434        "yuv420p",435        str(output_path),436    ]437    subprocess.run(command, check=True)438 439 440def resolve_raw_video(video_dir: Path, raw_video: Path | None) -> Path | None:441    if raw_video is not None and raw_video.exists():442        return raw_video443    video_id = video_dir.name444    raw_root = REPO_ROOT / "raw_video"445    for extension in VIDEO_EXTENSIONS:446        candidate = raw_root / f"{video_id}{extension}"447        if candidate.exists():448            return candidate449    return None450 451 452def extract_video_frames(raw_video: Path, fps: int, temp_dir: Path) -> List[Path]:453    temp_dir.mkdir(parents=True, exist_ok=True)454    command = [455        "ffmpeg",456        "-hide_banner",457        "-loglevel",458        "error",459        "-y",460        "-i",461        str(raw_video),462        "-vf",463        f"fps={fps}",464        str(temp_dir / "%08d.png"),465    ]466    subprocess.run(command, check=True)467    return sorted(temp_dir.glob("*.png"))468 469 470def render_overlay_frames(471    npz_paths: List[Path],472    raw_frame_paths: List[Path],473    overlay_dir: Path,474    draw_style: str,475    conf_threshold: float,476) -> None:477    overlay_dir.mkdir(parents=True, exist_ok=True)478    frame_count = min(len(npz_paths), len(raw_frame_paths))479    for index, (npz_path, raw_frame_path) in enumerate(zip(npz_paths[:frame_count], raw_frame_paths[:frame_count]), start=1):480        frame = load_npz_frame(npz_path, aggregated_index=index - 1 if npz_path.name == "poses.npz" else 0)481        pose_rgba = render_pose_image(frame, draw_style=draw_style, transparent=True, conf_threshold=conf_threshold)482        with Image.open(raw_frame_path) as raw_image:483            base = raw_image.convert("RGBA")484        overlay = Image.alpha_composite(base, pose_rgba)485        overlay.save(overlay_dir / f"{npz_path.stem}.png")486        if index == 1 or index % 100 == 0 or index == frame_count:487            print(f"Rendered overlay frame {index}/{frame_count}: {npz_path.name}")488 489 490def main() -> None:491    args = parse_args()492    args.draw_style = normalize_draw_style(args.draw_style)493    video_dir = args.video_dir.resolve()494    npz_dir = (args.npz_dir or (video_dir / "npz")).resolve()495    output_dir = (args.output_dir or (video_dir / f"visualization_{args.draw_style}")).resolve()496    pose_frame_dir = output_dir / "pose_frames"497    single_frame_dir = output_dir / "single_frames"498    overlay_frame_dir = output_dir / "overlay_frames"499    pose_video_path = output_dir / f"visualization_{args.draw_style}.mp4"500    overlay_video_path = output_dir / f"visualization_{args.draw_style}_overlay.mp4"501 502    if not npz_dir.exists():503        raise FileNotFoundError(f"NPZ directory not found: {npz_dir}")504 505    poses_npz_path = npz_dir / "poses.npz"506    if poses_npz_path.exists():507        npz_paths = [poses_npz_path]508    else:509        npz_paths = sorted(npz_dir.glob("*.npz"))510    if args.max_frames is not None:511        npz_paths = npz_paths[: args.max_frames]512    if not npz_paths:513        raise FileNotFoundError(f"No NPZ files found in {npz_dir}")514 515    if output_dir.exists() and args.force:516        shutil.rmtree(output_dir)517    output_dir.mkdir(parents=True, exist_ok=True)518 519    preview_indices = parse_frame_indices(args.frame_indices)520    preview_paths = [521        npz_paths[index - 1]522        for index in preview_indices523        if 0 < index <= len(npz_paths)524    ]525    save_frame_previews(preview_paths, single_frame_dir, args.draw_style, args.conf_threshold)526 527    render_pose_frames(npz_paths, pose_frame_dir, args.draw_style, args.conf_threshold)528    create_video_from_frames(pose_frame_dir, pose_video_path, args.fps)529 530    raw_video = resolve_raw_video(video_dir, args.raw_video)531    if raw_video is None:532        print("No raw video found for overlay rendering. Pose-only outputs were created.")533        return534 535    temp_root = Path(tempfile.mkdtemp(prefix="sign_dwpose_overlay_"))536    try:537        raw_frame_paths = extract_video_frames(raw_video, args.fps, temp_root)538        render_overlay_frames(npz_paths, raw_frame_paths, overlay_frame_dir, args.draw_style, args.conf_threshold)539        create_video_from_frames(overlay_frame_dir, overlay_video_path, args.fps)540    finally:541        shutil.rmtree(temp_root, ignore_errors=True)542 543 544if __name__ == "__main__":545    main()546