# ========================================================================= # 🎯 [路徑防禦補丁] 確保子程序被劫持進 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)