scripts/data_processors/smpl/generate_smpls.py

# =========================================================================
# 🎯 [路徑防禦補丁] 確保子程序被劫持進 venv 後，依然認得自己的專案根目錄路徑
# =========================================================================
import os
import sys
current_script_dir = os.path.dirname(os.path.abspath(__file__))
# 向上推 4 層，精準定位到 /app/champ-master 根目錄
master_root = os.path.abspath(os.path.join(current_script_dir, "..", "..", "..", ".."))
fourd_root = os.path.join(master_root, "4D-Humans")

if master_root not in sys.path:
    sys.path.insert(0, master_root)
if fourd_root not in sys.path:
    sys.path.insert(0, fourd_root)
# =========================================================================

import cv2
from pathlib import Path
import os
import argparse
import torch
import numpy as np
from tqdm import tqdm
import platform
import pyrender

from scripts.pretrained_models import (
    DETECTRON2_MODEL_PATH,
    HMR2_DEFAULT_CKPT,
)

if "PYOPENGL_PLATFORM" not in os.environ:
    os.environ["PYOPENGL_PLATFORM"] = "egl"

from hmr2.models import load_hmr2
from hmr2.utils import recursive_to
from hmr2.datasets.vitdet_dataset import ViTDetDataset
from hmr2.utils.renderer import Renderer, cam_crop_to_full
from .smpl_visualizer import SemanticRenderer

# For Windows, remove PYOPENGL_PLATFORM to enable default rendering backend
sys_name = platform.system()
if sys_name == "Windows":
    os.environ.pop("PYOPENGL_PLATFORM", None)

LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)

def predict_smpl(batch, model, model_cfg, figure_scale=None):
    all_verts = []
    all_cam_t = []
    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
        out = model(batch)

    pred_cam = out["pred_cam"]
    pred_smpl_parameter = out["pred_smpl_params"]
    if figure_scale is not None:
        pred_smpl_parameter['betas'][0][1] = float(figure_scale)
    smpl_output = model.smpl(
        **{k: v.float() for k, v in pred_smpl_parameter.items()},
        pose2rot=False,
    )
    pred_vertices = smpl_output.vertices
    out["pred_vertices"] = pred_vertices.reshape(
        batch["img"].shape[0], -1, 3
    )

    box_center = batch["box_center"].float()
    box_size = batch["box_size"].float()
    img_size = batch["img_size"].float()

    scaled_focal_length = (
        model_cfg.EXTRA.FOCAL_LENGTH
        / model_cfg.MODEL.IMAGE_SIZE
        * img_size.max()
    )
    pred_cam_t_full = cam_crop_to_full(
                pred_cam, box_center, box_size, img_size, scaled_focal_length
            ).detach().cpu().numpy()
    # Render the result
    batch_size = batch["img"].shape[0]
    for n in range(batch_size):
        # Add all verts and cams to list
        verts = out["pred_vertices"][n].detach().cpu().numpy()
        cam_t = pred_cam_t_full[n]
        all_verts.append(verts)
        all_cam_t.append(cam_t)

    misc_args = dict(
        mesh_base_color=LIGHT_BLUE,
        scene_bg_color=(1, 1, 1),
        focal_length=scaled_focal_length,
    )
    
    smpl_outs = {
        k: v.detach().cpu().numpy() for k, v in pred_smpl_parameter.items()
    }

    results_dict_for_rendering = {
        "verts": all_verts,
        "cam_t": all_cam_t,
        "render_res": img_size[n].cpu().numpy(),
        "smpls": smpl_outs,
        "scaled_focal_length": scaled_focal_length.cpu().numpy(),
    }
    return results_dict_for_rendering, misc_args

def load_image(img_cv2, detector):
    # Detect humans in image
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        det_out = detector(img_cv2)
    det_instances = det_out["instances"]
    valid_idx = (det_instances.pred_classes == 0) & (det_instances.scores > 0.5)
    boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
    # Run HMR2.0 on all detected humans
    dataset = ViTDetDataset(model_cfg, img_cv2, boxes)
    return torch.utils.data.DataLoader(
        dataset, batch_size=8, shuffle=False, num_workers=0
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Inference SMPL with 4D-Humans")
    parser.add_argument("--device", type=int, default=0, help="GPU device ID")
    parser.add_argument(
        "--reference_imgs_folder",
        type=str,
        default="",
        help="Folder path to reference imgs",
    )
    parser.add_argument(
        "--driving_video_path",
        type=str,
        default="driving_videos",
        help="Folder path to driving videos",
    )
    parser.add_argument(
        "--figure_scale",
        type=int,
        default=None,
        help="Adjust the figure scale to better fit extreme shape",
    )

    args = parser.parse_args()

    reference_img_paths = []
    if args.reference_imgs_folder:
        os.makedirs(args.reference_imgs_folder, exist_ok=True)
        os.makedirs(
            os.path.join(args.reference_imgs_folder, "visualized_imgs"), exist_ok=True
        )
        os.makedirs(os.path.join(args.reference_imgs_folder, "mask"), exist_ok=True)
        os.makedirs(
            os.path.join(args.reference_imgs_folder, "semantic_map"), exist_ok=True
        )
        os.makedirs(os.path.join(args.reference_imgs_folder, "depth"), exist_ok=True)
        os.makedirs(
            os.path.join(args.reference_imgs_folder, "smpl_results"), exist_ok=True
        )

        reference_img_paths = [
            path
            for path in os.listdir(os.path.join(args.reference_imgs_folder, "images"))
        ]

    driving_videos_paths = [args.driving_video_path]

    model, model_cfg = load_hmr2(HMR2_DEFAULT_CKPT)

    # Load ultra-lightweight torchvision detector to completely eliminate ViT-H overhead
    import torchvision
    class LightweightDetector:
        def __init__(self, device):
            print("Loading Faster R-CNN ResNet-50 detector (High Accuracy & Low VRAM)...")
            self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn_v2(weights=torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
            self.model.to(device)
            self.model.eval()
            self.device = device
            
        def __call__(self, img_cv2):
            img_rgb = img_cv2[:, :, ::-1].copy()
            img_tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float() / 255.0
            img_tensor = img_tensor.unsqueeze(0).to(self.device)
            
            with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
                prediction = self.model(img_tensor)[0]
                
            labels = prediction['labels']
            scores = prediction['scores']
            boxes = prediction['boxes']
            
            # Person class is 1 in COCO
            valid_idx = (labels == 1) & (scores > 0.5)
            person_boxes = boxes[valid_idx].cpu().numpy()
            person_scores = scores[valid_idx].cpu().numpy()
            
            # Protection: Only take the person with the highest confidence score
            if len(person_boxes) > 1:
                best_idx = np.argmax(person_scores)
                person_boxes = person_boxes[best_idx:best_idx+1]
                person_scores = person_scores[best_idx:best_idx+1]
            
            class MockInstances:
                def __init__(self, boxes, scores):
                    class MockBoxes:
                        def __init__(self, b):
                            self.tensor = torch.from_numpy(b)
                    self.pred_boxes = MockBoxes(boxes)
                    self.scores = scores
                    self.pred_classes = np.zeros(len(boxes), dtype=int)
            
            return {"instances": MockInstances(person_boxes, person_scores)}

    detector = LightweightDetector(args.device)

    model = model.to(args.device)
    detector.model.to(args.device)
    # This PyRender is only used for visualizing, we use Blender after to render different conditions.
    renderer = SemanticRenderer(
        model_cfg,
        faces=model.smpl.faces,
        lbs=model.smpl.lbs_weights,
        viewport_size=(720, 720),
    )
    for img_path in tqdm(reference_img_paths, desc="Processing Reference Images:"):
        img_fn, _ = os.path.splitext(os.path.basename(img_path))
        if os.path.exists(os.path.join(args.reference_imgs_folder, "smpl_results", f"{img_fn}.npy")):
            continue

        img_cv2 = cv2.imread(
            str(os.path.join(args.reference_imgs_folder, "images", img_path))
        )

        try:
            renderer.renderer.delete()
        except Exception:
            pass

        import torch
        torch.cuda.empty_cache()
        print(f"\n[DEBUG] Creating OffscreenRenderer for {img_path} ({img_cv2.shape})", flush=True)

        renderer.renderer = pyrender.OffscreenRenderer(
            viewport_width=img_cv2.shape[:2][::-1][0],
            viewport_height=img_cv2.shape[:2][::-1][1],
            point_size=1.0,
        )
        img_fn, _ = os.path.splitext(os.path.basename(img_path))
        dataloader = load_image(img_cv2, detector)

        for batch in dataloader:
            batch = recursive_to(batch, args.device)
            results_dict_for_rendering, misc_args = predict_smpl(batch, model, model_cfg, args.figure_scale)
            
            print(f"[DEBUG] Rendering {img_path} with PyRender...", flush=True)
            rendering_results = renderer.render_all_multiple(
                results_dict_for_rendering["verts"], 
                cam_t=results_dict_for_rendering["cam_t"], 
                render_res=results_dict_for_rendering["render_res"], **misc_args
            )
            print(f"[DEBUG] Render finished for {img_path}", flush=True)
            # Overlay image
            valid_mask = rendering_results["Image"][:, :, -1][:, :, np.newaxis]
            cam_view = (
                valid_mask * rendering_results["Image"][:, :, [2, 1, 0]]
                + (1 - valid_mask) * img_cv2.astype(np.float32)[:, :, ::-1] / 255
            )

            cv2.imwrite(
                os.path.join(args.reference_imgs_folder, "visualized_imgs", f"{img_fn}.png"),
                255 * cam_view[:, :, ::-1])
            cv2.imwrite(
                os.path.join(args.reference_imgs_folder, "mask", f"{img_fn}.png"),
                255 * rendering_results.get("Mask")[:, :, 0])
            cv2.imwrite(
                os.path.join(args.reference_imgs_folder, "semantic_map", f"{img_fn}.png"),
                255 * rendering_results.get("SemanticMap"))
            np.save(
                str(os.path.join(args.reference_imgs_folder, "smpl_results", f"{img_fn}.npy")),
                results_dict_for_rendering)
    
    for video_path in tqdm(driving_videos_paths, desc="Processing Driving Videos:"):
        os.makedirs(video_path, exist_ok=True)
        os.makedirs(os.path.join(video_path, "smpl_results"), exist_ok=True)

        driving_img_paths = [
            path for path in os.listdir(os.path.join(video_path, "images"))
        ]
        driving_img_paths.sort(key=lambda x: int(x.split(".")[0]))
        smpls = []
        cams = []
        # ---------------------------------------------------------
        # Phase 1: Detection (Run Detectron2 on all frames to avoid VRAM thrashing)
        # ---------------------------------------------------------
        print(f"[{video_path}] Phase 1: Running Person Detection...", flush=True)
        # 釋放 HMR2 到 CPU，把 VRAM 全留給 Detectron2
        model = model.cpu()
        torch.cuda.empty_cache()
        detector.model = detector.model.to(args.device)

        frame_boxes = {}
        for img_path in tqdm(driving_img_paths, desc="Phase 1 (Detection)"):
            img_cv2 = cv2.imread(str(os.path.join(video_path, "images", img_path)))
            
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                det_out = detector(img_cv2)
                
            det_instances = det_out["instances"]
            valid_idx = (det_instances.pred_classes == 0) & (det_instances.scores > 0.5)
            boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
            frame_boxes[img_path] = boxes
            
        # ---------------------------------------------------------
        # Phase 2: HMR2 (Run pose estimation on detected boxes)
        # ---------------------------------------------------------
        print(f"[{video_path}] Phase 2: Running 3D Pose Estimation...", flush=True)
        # 釋放 Detectron2 到 CPU，把 VRAM 全留給 HMR2
        detector.model = detector.model.cpu()
        torch.cuda.empty_cache()
        model = model.to(args.device)

        smpls = []
        cams = []
        for img_path in tqdm(driving_img_paths, desc="Phase 2 (HMR2)"):
            img_cv2 = cv2.imread(str(os.path.join(video_path, "images", img_path)))
            img_fn, _ = os.path.splitext(os.path.basename(img_path))
            boxes = frame_boxes[img_path]
            
            dataset = ViTDetDataset(model_cfg, img_cv2, boxes)
            dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0)

            for batch in dataloader:
                batch = recursive_to(batch, args.device)
                results_dict_for_rendering, misc_args = predict_smpl(batch, model, model_cfg)
                cams.append(results_dict_for_rendering["cam_t"][0])
                smpls.append(results_dict_for_rendering["smpls"])
                np.save(
                    str(os.path.join(video_path, "smpl_results", f"{img_fn}.npy")),
                    results_dict_for_rendering)
                
        np.savez(
            str(os.path.join(video_path, "smpl_results", f"smpls_group.npz")),
            smpl=smpls,camera=cams)