roi/EditP23

1import os2import argparse3import numpy as np4import torch5from PIL import Image6from torchvision.transforms import v27from omegaconf import OmegaConf8from einops import rearrange9from tqdm import tqdm10from huggingface_hub import hf_hub_download11import sys12 13script_dir = os.path.dirname(os.path.abspath(__file__))14submodule_path = os.path.join(script_dir, "..", "external", "instant-mesh")15sys.path.insert(0, submodule_path)16 17from src.utils.camera_util import (18    get_circular_camera_poses,19    get_zero123plus_input_cameras,20    FOV_to_intrinsics,21)22from src.utils.train_util import instantiate_from_config23from src.utils.mesh_util import save_obj24from src.utils.infer_util import save_video25 26 27def get_render_cameras(28    batch_size=1, M=120, radius=4.0, elevation=20.0, is_flexicubes=False29):30    c2ws = get_circular_camera_poses(M=M, radius=radius, elevation=elevation)31    if is_flexicubes:32        cameras = torch.linalg.inv(c2ws)33        cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1, 1)34    else:35        extrinsics = c2ws.flatten(-2)36        intrinsics = (37            FOV_to_intrinsics(30.0).unsqueeze(0).repeat(M, 1, 1).float().flatten(-2)38        )39        cameras = torch.cat([extrinsics, intrinsics], dim=-1)40        cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1)41    return cameras42 43 44def render_frames(45    model, planes, render_cameras, render_size=512, chunk_size=1, is_flexicubes=False46):47    frames = []48    for i in tqdm(range(0, render_cameras.shape[1], chunk_size)):49        if is_flexicubes:50            frame = model.forward_geometry(51                planes, render_cameras[:, i : i + chunk_size], render_size=render_size52            )["img"]53        else:54            frame = model.forward_synthesizer(55                planes, render_cameras[:, i : i + chunk_size], render_size=render_size56            )["images_rgb"]57        frames.append(frame)58    frames = torch.cat(frames, dim=1)[0]59    return frames60 61def main(args):62    """63    Main function to run the 3D mesh generation process.64    """65    # ============================66    # CONFIG67    # ============================68    print("🚀 Starting 3D mesh generation...")69    config = OmegaConf.load(args.config)70    config_name = os.path.basename(args.config).replace(".yaml", "")71    model_config = config.model_config72    infer_config = config.infer_config73    IS_FLEXICUBES = config_name.startswith("instant-mesh")74 75    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")76    print(f"Using device: {device}")77 78    # ============================79    # SETUP OUTPUT DIRECTORY80    # ============================81    os.makedirs(args.output_dir, exist_ok=True)82    base_name = os.path.splitext(os.path.basename(args.input_file))[0]83    mesh_path = os.path.join(args.output_dir, "recon.obj")84    video_path = os.path.join(args.output_dir, "recon.mp4")85 86    # ============================87    # LOAD RECONSTRUCTION MODEL88    # ============================89    print("Loading reconstruction model...")90    model = instantiate_from_config(model_config)91 92    # Download model checkpoint if it doesn't exist93    model_ckpt_path = (94        infer_config.model_path95        if os.path.exists(infer_config.model_path)96        else hf_hub_download(97            repo_id="TencentARC/InstantMesh",98            filename=f"{config_name.replace('-', '_')}.ckpt",99            repo_type="model",100        )101    )102 103    # Load the state dictionary104    state_dict = torch.load(model_ckpt_path, map_location="cpu")["state_dict"]105    state_dict = {106        k[14:]: v for k, v in state_dict.items() if k.startswith("lrm_generator.")107    }108    model.load_state_dict(state_dict, strict=True)109    model = model.to(device).eval()110 111    if IS_FLEXICUBES:112        model.init_flexicubes_geometry(device, fovy=30.0)113 114    # ============================115    # PREPARE DATA116    # ============================117    print(f"Processing input file: {args.input_file}")118    119    # Load and preprocess the input image120    input_image = Image.open(args.input_file).convert("RGB")121    images = np.asarray(input_image, dtype=np.float32) / 255.0122    images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float()123    # Rearrange from (C, H, W) to (B, C, H, W) where B is the number of views124    images = rearrange(images, "c (n h) (m w) -> (n m) c h w", n=3, m=2)125    images = images.unsqueeze(0).to(device)126    images = v2.functional.resize(images, size=320, interpolation=3, antialias=True).clamp(0, 1)127 128    input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0 * args.scale).to(device)129 130    # ============================131    # RUN INFERENCE AND SAVE OUTPUT132    # ============================133    with torch.no_grad():134        # Generate 3D mesh135        planes = model.forward_planes(images, input_cameras)136        mesh_out = model.extract_mesh(planes, use_texture_map=False, **infer_config)137        138        # Save the mesh139        vertices, faces, vertex_colors = mesh_out140        save_obj(vertices, faces, vertex_colors, mesh_path)141        print(f"✅ Mesh saved to {mesh_path}")142 143        # Render and save video if enabled144        if args.save_video:145            print("🎥 Rendering video...")146            render_size = infer_config.render_resolution147            chunk_size = 20 if IS_FLEXICUBES else 1148            render_cameras = get_render_cameras(149                batch_size=1,150                M=120,151                radius=args.distance,152                elevation=20.0,153                is_flexicubes=IS_FLEXICUBES,154            ).to(device)155 156            frames = render_frames(157                model=model,158                planes=planes,159                render_cameras=render_cameras,160                render_size=render_size,161                chunk_size=chunk_size,162                is_flexicubes=IS_FLEXICUBES,163            )164            save_video(frames, video_path, fps=30)165            print(f"✅ Video saved to {video_path}")166 167    print("✨ Process complete.")168 169if __name__ == "__main__":170    # ============================171    # SCRIPT ARGUMENTS172    # ============================173    parser = argparse.ArgumentParser(174        description="Generate a 3D mesh and video from a single multi-view PNG file."175    )176    177    # Positional argument for config file178    parser.add_argument(179        "config", 180        type=str, 181        help="Path to the model config file (.yaml)."182    )183    184    # Required file paths185    parser.add_argument(186        "--input_file", 187        type=str, 188        required=True, 189        help="Path to the input PNG file."190    )191    parser.add_argument(192        "--output_dir", 193        type=str, 194        default="outputs/", 195        help="Directory to save the output .obj and .mp4 files. Defaults to 'outputs/'."196    )197    198    # Optional parameters for model and rendering199    parser.add_argument(200        "--scale", 201        type=float, 202        default=1.0, 203        help="Scale of the input cameras."204    )205    parser.add_argument(206        "--distance", 207        type=float, 208        default=4.5, 209        help="Camera distance for rendering the output video."210    )211    parser.add_argument(212        "--no_video", 213        dest="save_video", 214        action="store_false", 215        help="If set, disables saving the output .mp4 video."216    )217    218    parsed_args = parser.parse_args()219    main(parsed_args)220