import os, subprocess, shlex, sys, gc import time import torch import numpy as np import shutil import argparse import gradio as gr import uuid import spaces subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl")) subprocess.run(shlex.split("pip install wheel/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl")) subprocess.run(shlex.split("pip install wheel/curope-0.0.0-cp310-cp310-linux_x86_64.whl")) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) os.sys.path.append(os.path.abspath(os.path.join(BASE_DIR, "submodules", "dust3r"))) # os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' from dust3r.inference import inference from dust3r.model import AsymmetricCroCo3DStereo from dust3r.utils.device import to_numpy from dust3r.image_pairs import make_pairs from dust3r.cloud_opt import global_aligner, GlobalAlignerMode from utils.dust3r_utils import compute_global_alignment, load_images, storePly, save_colmap_cameras, save_colmap_images from argparse import ArgumentParser, Namespace from arguments import ModelParams, PipelineParams, OptimizationParams from train_joint import training from render_by_interp import render_sets GRADIO_CACHE_FOLDER = './gradio_cache_folder' ############################################################################################################################################# def get_dust3r_args_parser(): parser = argparse.ArgumentParser() parser.add_argument("--image_size", type=int, default=512, choices=[512, 224], help="image size") parser.add_argument("--model_path", type=str, default="submodules/dust3r/checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth", help="path to the model weights") parser.add_argument("--device", type=str, default='cuda', help="pytorch device") parser.add_argument("--batch_size", type=int, default=1) parser.add_argument("--schedule", type=str, default='linear') parser.add_argument("--lr", type=float, default=0.01) parser.add_argument("--niter", type=int, default=300) parser.add_argument("--focal_avg", type=bool, default=True) parser.add_argument("--n_views", type=int, default=3) parser.add_argument("--base_path", type=str, default=GRADIO_CACHE_FOLDER) return parser @spaces.GPU(duration=150) def process(inputfiles, input_path=None): if input_path is not None: imgs_path = './assets/example/' + input_path imgs_names = sorted(os.listdir(imgs_path)) inputfiles = [] for imgs_name in imgs_names: file_path = os.path.join(imgs_path, imgs_name) print(file_path) inputfiles.append(file_path) print(inputfiles) # ------ (1) Coarse Geometric Initialization ------ # os.system(f"rm -rf {GRADIO_CACHE_FOLDER}") parser = get_dust3r_args_parser() opt = parser.parse_args() tmp_user_folder = str(uuid.uuid4()).replace("-", "") opt.img_base_path = os.path.join(opt.base_path, tmp_user_folder) img_folder_path = os.path.join(opt.img_base_path, "images") img_folder_path = os.path.join(opt.img_base_path, "images") model = AsymmetricCroCo3DStereo.from_pretrained(opt.model_path).to(opt.device) os.makedirs(img_folder_path, exist_ok=True) opt.n_views = len(inputfiles) if opt.n_views == 1: raise gr.Error("The number of input images should be greater than 1.") print("Multiple images: ", inputfiles) for image_path in inputfiles: if input_path is not None: shutil.copy(image_path, img_folder_path) else: shutil.move(image_path, img_folder_path) train_img_list = sorted(os.listdir(img_folder_path)) assert len(train_img_list)==opt.n_views, f"Number of images in the folder is not equal to {opt.n_views}" images, ori_size, imgs_resolution = load_images(img_folder_path, size=512) resolutions_are_equal = len(set(imgs_resolution)) == 1 if resolutions_are_equal == False: raise gr.Error("The resolution of the input image should be the same.") print("ori_size", ori_size) start_time = time.time() ###################################################### pairs = make_pairs(images, scene_graph='complete', prefilter=None, symmetrize=True) output = inference(pairs, model, opt.device, batch_size=opt.batch_size) output_colmap_path=img_folder_path.replace("images", "sparse/0") os.makedirs(output_colmap_path, exist_ok=True) scene = global_aligner(output, device=opt.device, mode=GlobalAlignerMode.PointCloudOptimizer) loss = compute_global_alignment(scene=scene, init="mst", niter=opt.niter, schedule=opt.schedule, lr=opt.lr, focal_avg=opt.focal_avg) scene = scene.clean_pointcloud() imgs = to_numpy(scene.imgs) focals = scene.get_focals() poses = to_numpy(scene.get_im_poses()) pts3d = to_numpy(scene.get_pts3d()) scene.min_conf_thr = float(scene.conf_trf(torch.tensor(1.0))) confidence_masks = to_numpy(scene.get_masks()) intrinsics = to_numpy(scene.get_intrinsics()) ###################################################### end_time = time.time() print(f"Time taken for {opt.n_views} views: {end_time-start_time} seconds") save_colmap_cameras(ori_size, intrinsics, os.path.join(output_colmap_path, 'cameras.txt')) save_colmap_images(poses, os.path.join(output_colmap_path, 'images.txt'), train_img_list) pts_4_3dgs = np.concatenate([p[m] for p, m in zip(pts3d, confidence_masks)]) color_4_3dgs = np.concatenate([p[m] for p, m in zip(imgs, confidence_masks)]) color_4_3dgs = (color_4_3dgs * 255.0).astype(np.uint8) storePly(os.path.join(output_colmap_path, "points3D.ply"), pts_4_3dgs, color_4_3dgs) pts_4_3dgs_all = np.array(pts3d).reshape(-1, 3) np.save(output_colmap_path + "/pts_4_3dgs_all.npy", pts_4_3dgs_all) np.save(output_colmap_path + "/focal.npy", np.array(focals.cpu())) ### save VRAM del scene torch.cuda.empty_cache() gc.collect() ################################################################################################################################################## # ------ (2) Fast 3D-Gaussian Optimization ------ parser = ArgumentParser(description="Training script parameters") lp = ModelParams(parser) op = OptimizationParams(parser) pp = PipelineParams(parser) parser.add_argument('--debug_from', type=int, default=-1) parser.add_argument("--test_iterations", nargs="+", type=int, default=[]) parser.add_argument("--save_iterations", nargs="+", type=int, default=[]) parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[]) parser.add_argument("--start_checkpoint", type=str, default = None) parser.add_argument("--scene", type=str, default="demo") parser.add_argument("--n_views", type=int, default=3) parser.add_argument("--get_video", action="store_true") parser.add_argument("--optim_pose", type=bool, default=True) parser.add_argument("--skip_train", action="store_true") parser.add_argument("--skip_test", action="store_true") args = parser.parse_args(sys.argv[1:]) args.save_iterations.append(args.iterations) args.model_path = opt.img_base_path + '/output/' args.source_path = opt.img_base_path # args.model_path = GRADIO_CACHE_FOLDER + '/output/' # args.source_path = GRADIO_CACHE_FOLDER args.iteration = 1000 os.makedirs(args.model_path, exist_ok=True) training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from, args) ################################################################################################################################################## # ------ (3) Render video by interpolation ------ parser = ArgumentParser(description="Testing script parameters") model = ModelParams(parser, sentinel=True) pipeline = PipelineParams(parser) args.eval = True args.get_video = True args.n_views = opt.n_views render_sets( model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test, args, ) output_ply_path = opt.img_base_path + f'/output/point_cloud/iteration_{args.iteration}/point_cloud.ply' output_video_path = opt.img_base_path + f'/output/demo_{opt.n_views}_view.mp4' # output_ply_path = GRADIO_CACHE_FOLDER+ f'/output/point_cloud/iteration_{args.iteration}/point_cloud.ply' # output_video_path = GRADIO_CACHE_FOLDER+ f'/output/demo_{opt.n_views}_view.mp4' return output_video_path, output_ply_path, output_ply_path ################################################################################################################################################## _TITLE = '''InstantSplat''' _DESCRIPTION = '''
InstantSplat: Sparse-view SfM-free Gaussian Splatting in Seconds

    arxiv

* Official demo of: [InstantSplat: Sparse-view SfM-free Gaussian Splatting in Seconds](https://instantsplat.github.io/). * Sparse-view examples for direct viewing: you can simply click the examples (in the bottom of the page), to quickly view the results on representative data. * Training speeds may slow if the resolution or number of images is large. To achieve performance comparable to what has been reported, please conduct tests on your own GPU (A100/4090). ''' # Github Source Code  #   # Nvidia # * If InstantSplat is helpful, please give us a star ⭐ on Github. Thanks! # block = gr.Blocks(title=_TITLE).queue() block = gr.Blocks().queue() with block: with gr.Row(): with gr.Column(scale=1): # gr.Markdown('# ' + _TITLE) gr.Markdown(_DESCRIPTION) with gr.Row(variant='panel'): with gr.Tab("Input"): inputfiles = gr.File(file_count="multiple", label="images") input_path = gr.Textbox(visible=False, label="example_path") button_gen = gr.Button("RUN") with gr.Row(variant='panel'): with gr.Tab("Output"): with gr.Column(scale=2): output_model = gr.Model3D( label="3D Model (Gaussian)", # height=300, interactive=False, # clear_color=[1.0, 1.0, 1.0, 1.0] ) output_file = gr.File(label="ply") with gr.Column(scale=1): output_video = gr.Video(label="video") button_gen.click(process, inputs=[inputfiles], outputs=[ output_video, output_file, output_model]) gr.Examples( examples=[ "sora-santorini-3-views", "TT-family-3-views", "dl3dv-ba55-3-views", ], inputs=[input_path], outputs=[output_video, output_file, output_model], fn=lambda x: process(inputfiles=None, input_path=x), cache_examples=True, label='Sparse-view Examples' ) block.launch(server_name="0.0.0.0", share=False)