aggregate all scripts into one function

.gitignore

@@ -23,3 +23,4 @@ build/
 /dev/
 gradio_cache_folder/
 results/
+.gradio

app.py

@@ -10,12 +10,18 @@ import re
 import torch
 import spaces
 
-subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl --force-reinstall"))
-subprocess.run(shlex.split("pip install wheel/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl --force-reinstall"))
-subprocess.run(shlex.split("pip install wheel/curope-0.0.0-cp310-cp310-linux_x86_64.whl --force-reinstall"))
-
+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, "dynamic_predictor")))
+os.sys.path.append(os.path.abspath(os.path.join(BASE_DIR)))
 GRADIO_CACHE_FOLDER = './gradio_cache_folder'
 
+from dynamic_predictor.launch import main as dynamic_predictor_main
+from utils_das3r.rearrange import main as rearrange_main
+from train_gui import main as train_main
+from render import main as render_main
 
 def get_dust3r_args_parser():
     parser = argparse.ArgumentParser()
@@ -37,25 +43,7 @@ def natural_sort(l):
     alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key.split('/')[-1])]
     return sorted(l, key=alphanum_key)
 
-def cmd(command):
-    print(command)
-    subprocess.run(shlex.split(command))
-
-@spaces.GPU(duration=70)
-def cmd_gpu_s1(command):
-    print('gpu:', command)
-    subprocess.run(shlex.split(command))
-
-@spaces.GPU(duration=40)
-def cmd_gpu_s2(command):
-    print('gpu:', command)
-    subprocess.run(shlex.split(command))
-
-@spaces.GPU(duration=20)
-def cmd_gpu_s3(command):
-    print('gpu:', command)
-    subprocess.run(shlex.split(command))
-
+@spaces.GPU(duration=150)
 def process(inputfiles, input_path='demo'):
     if inputfiles:
         frames = natural_sort(inputfiles)
@@ -75,22 +63,17 @@ def process(inputfiles, input_path='demo'):
 
     imgs_path = temp_dir
     output_path = f'./results/{input_path}/output'
-    cmd_gpu_s1(f"python dynamic_predictor/launch.py --mode=eval_pose_custom \
-        --pretrained=Kai422kx/das3r \
-        --dir_path={imgs_path} \
-        --output_dir={output_path} \
-        --use_pred_mask --n_iter 150")
-    
-    cmd(f"python utils/rearrange.py --output_dir={output_path}")
-    output_path = f'{output_path}_rearranged'
+    rearranged_path = f'{output_path}_rearranged'
 
+    dynamic_predictor_main(pretrained='Kai422kx/das3r', dir_path=imgs_path, output_dir=output_path, use_pred_mask=True, n_iter=150)
+    rearrange_main(output_dir=output_path, rearranged_path = rearranged_path)
+    train_main(s = rearranged_path, m = rearranged_path, iter = 2000)
+    render_main(s = rearranged_path, m = rearranged_path, iter = 2000, get_video = True)
 
-    cmd_gpu_s2(f"python train_gui.py -s {output_path} -m {output_path} --iter 2000")
-    cmd_gpu_s3(f"python render.py -s {output_path} -m {output_path} --iter 2000 --get_video")
 
 
-    output_video_path = f"{output_path}/rendered.mp4"
-    output_ply_path = f"{output_path}/point_cloud/iteration_2000/point_cloud.ply"
+    output_video_path = f"{rearranged_path}/rendered.mp4"
+    output_ply_path = f"{rearranged_path}/point_cloud/iteration_2000/point_cloud.ply"
     return  output_video_path, output_ply_path, output_ply_path
 
 

dynamic_predictor/dust3r/training.py

@@ -217,6 +217,7 @@ def train(args):
     data_loader_test = {}
     for dataset in args.test_dataset.split('+'):
         testset = build_dataset(dataset, args.test_batch_size, args.num_workers, test=True)
+        print(args.test_dataset)
         name_testset = dataset.split('(')[0]
         if getattr(testset.dataset.dataset, 'strides', None) is not None:
             name_testset += f'_stride{testset.dataset.dataset.strides}'

dynamic_predictor/launch.py

@@ -10,6 +10,35 @@ import torch.backends.cudnn as cudnn
 import numpy as np
 import os
 
+def main(pretrained, dir_path, output_dir, use_pred_mask, n_iter):
+    
+    args = get_args_parser()
+    args = args.parse_args()
+    args.pretrained = pretrained
+    args.dir_path = dir_path
+    args.output_dir = output_dir
+    args.use_pred_mask = use_pred_mask
+    args.n_iter = n_iter
+
+    misc.init_distributed_mode(args)
+    global_rank = misc.get_rank()
+    world_size = misc.get_world_size()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(device)
+
+    # fix the seed
+    seed = args.seed + misc.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = args.cudnn_benchmark
+    model, _ = load_model(args, device)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+
+    pose_estimation_custom(args, model, device, save_dir=args.output_dir)
+
+
+
 if __name__ == '__main__':
     args = get_args_parser()
     args = args.parse_args()
@@ -39,3 +68,4 @@ if __name__ == '__main__':
 
         exit(0)
     train(args)
+

gaussian_renderer/__init__.py

@@ -16,8 +16,8 @@ from diff_gaussian_rasterization import (
     GaussianRasterizer,
 )
 from scene.gaussian_model import GaussianModel
-from utils.sh_utils import eval_sh
-from utils.pose_utils import get_camera_from_tensor, quadmultiply
+from utils_das3r.sh_utils import eval_sh
+from utils_das3r.pose_utils import get_camera_from_tensor, quadmultiply
 
 
 def render(

render.py

@@ -10,19 +10,20 @@
 #
 
 import torch
+import sys
 from scene import Scene
 import os
 from tqdm import tqdm
 from os import makedirs
 from gaussian_renderer import render_test as render
 import torchvision
-from utils.general_utils import safe_state
+from utils_das3r.general_utils import safe_state
 from argparse import ArgumentParser
 from arguments import ModelParams, PipelineParams, get_combined_args
 from gaussian_renderer import GaussianModel
-from utils.pose_utils import get_tensor_from_camera
-from utils.camera_utils import generate_interpolated_path
-from utils.camera_utils import visualizer
+from utils_das3r.pose_utils import get_tensor_from_camera
+from utils_das3r.camera_utils import generate_interpolated_path
+from utils_das3r.camera_utils import visualizer
 import cv2
 import numpy as np
 import imageio
@@ -122,6 +123,37 @@ def render_sets(
         output_video_file = os.path.join(dataset.model_path, f'rendered.mp4')
         images_to_video(image_folder, output_video_file, fps=15)
 
+def main(s, m, iter, get_video):
+    # Set up command line argument parser
+    parser = ArgumentParser(description="script parameters")
+    lp = ModelParams(parser)
+    pp = PipelineParams(parser)
+    parser.add_argument("--iteration", default=-1, type=int)
+    parser.add_argument("--skip_train", action="store_true")
+    parser.add_argument("--skip_test", action="store_true")
+    parser.add_argument("--quiet", action="store_true")
+
+    parser.add_argument("--get_video", action="store_true")
+    parser.add_argument("--n_views", default=None, type=int)
+    parser.add_argument("--scene", default=None, type=str)
+    
+    args = parser.parse_args(sys.argv[1:])
+
+    args.source_path = s
+    args.model_path = m
+    args.iteration = iter
+    args.get_video = get_video
+    # Initialize system state (RNG)
+    # safe_state(args.quiet)
+    args.eval = False
+    render_sets(
+        lp.extract(args),
+        args.iteration,
+        pp.extract(args),
+        args.skip_train,
+        args.skip_test,
+        args,
+    )
 
 if __name__ == "__main__":
     # Set up command line argument parser

scene/__init__.py

@@ -12,11 +12,11 @@
 import os
 import random
 import json
-from utils.system_utils import searchForMaxIteration
+from utils_das3r.system_utils import searchForMaxIteration
 from scene.dataset_readers import sceneLoadTypeCallbacks
 from scene.gaussian_model import GaussianModel
 from arguments import ModelParams
-from utils.camera_utils import cameraList_from_camInfos, camera_to_JSON
+from utils_das3r.camera_utils import cameraList_from_camInfos, camera_to_JSON
 import open3d as o3d
 
 class Scene:

scene/cameras.py

@@ -12,7 +12,7 @@
 import torch
 from torch import nn
 import numpy as np
-from utils.graphics_utils import getWorld2View2, getProjectionMatrix
+from utils_das3r.graphics_utils import getWorld2View2, getProjectionMatrix
 
 class Camera(nn.Module):
     def __init__(self, colmap_id, intr, R, T, original_pose, FoVx, FoVy, image, gt_alpha_mask, dynamic_mask, enlarged_dynamic_mask,

scene/dataset_readers.py

@@ -15,16 +15,16 @@ from PIL import Image
 from typing import NamedTuple
 from scene.colmap_loader import read_extrinsics_text, read_intrinsics_text, qvec2rotmat, \
     read_extrinsics_binary, read_intrinsics_binary, read_points3D_binary, read_points3D_text
-from utils.graphics_utils import getWorld2View2, focal2fov, fov2focal, to_open3d_point_cloud
+from utils_das3r.graphics_utils import getWorld2View2, focal2fov, fov2focal, to_open3d_point_cloud
 import numpy as np
 import json
 from pathlib import Path
 from plyfile import PlyData, PlyElement
-from utils.sh_utils import SH2RGB
+from utils_das3r.sh_utils import SH2RGB
 from scene.gaussian_model import BasicPointCloud
-from utils.vo_eval import file_interface
+from utils_das3r.vo_eval import file_interface
 import torch
-from utils.pose_utils import quad2rotation
+from utils_das3r.pose_utils import quad2rotation
 
 class CameraInfo(NamedTuple):
     uid: int

scene/gaussian_model.py

@@ -12,19 +12,19 @@
 import torch
 # from lietorch import SO3, SE3, Sim3, LieGroupParameter
 import numpy as np
-from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from utils_das3r.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
 from torch import nn
 import os
-from utils.system_utils import mkdir_p
+from utils_das3r.system_utils import mkdir_p
 from plyfile import PlyData, PlyElement
-from utils.sh_utils import RGB2SH
+from utils_das3r.sh_utils import RGB2SH
 from simple_knn._C import distCUDA2
-from utils.graphics_utils import BasicPointCloud
-from utils.general_utils import strip_symmetric, build_scaling_rotation
+from utils_das3r.graphics_utils import BasicPointCloud
+from utils_das3r.general_utils import strip_symmetric, build_scaling_rotation
 from scipy.spatial.transform import Rotation as R
-from utils.pose_utils import rotation2quad, get_tensor_from_camera
-from utils.graphics_utils import getWorld2View2
-from utils.pose_utils import rotation2quad, get_tensor_from_camera, depth_to_pts3d
+from utils_das3r.pose_utils import rotation2quad, get_tensor_from_camera
+from utils_das3r.graphics_utils import getWorld2View2
+from utils_das3r.pose_utils import rotation2quad, get_tensor_from_camera, depth_to_pts3d
 
 class GaussianModel:
 

train_gui.py

@@ -14,19 +14,19 @@ import numpy as np
 import torch
 from PIL import Image
 from random import randint
-from utils.loss_utils import l1_loss, ssim
+from utils_das3r.loss_utils import l1_loss, ssim
 from gaussian_renderer import render, network_gui, render_confidence
 import sys
 from scene import Scene, GaussianModel
-from utils.general_utils import safe_state
+from utils_das3r.general_utils import safe_state
 import uuid
 from tqdm import tqdm
-from utils.image_utils import psnr
+from utils_das3r.image_utils import psnr
 from argparse import ArgumentParser, Namespace
 from arguments import ModelParams, PipelineParams, OptimizationParams
-from utils.pose_utils import get_camera_from_tensor
-from utils.vo_eval import load_traj, eval_metrics, plot_trajectory
-from utils.gui_utils import orbit_camera, OrbitCamera
+from utils_das3r.pose_utils import get_camera_from_tensor
+from utils_das3r.vo_eval import load_traj, eval_metrics, plot_trajectory
+from utils_das3r.gui_utils import orbit_camera, OrbitCamera
 import dearpygui.dearpygui as dpg
 from scipy.spatial.transform import Rotation
 
@@ -710,7 +710,59 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i
                 tb_writer.add_scalar(config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration)
 
     return log
+
+def main(s, m, iter = 2000):
+    # Set up command line argument parser
+    parser = ArgumentParser(description="Training script parameters")
+    lp = ModelParams(parser)
+    op = OptimizationParams(parser)
+    pp = PipelineParams(parser)
+    parser.add_argument('--ip', type=str, default="127.0.0.1")
+    parser.add_argument('--port', type=int, default=6009)
+    parser.add_argument('--debug_from', type=int, default=-1)
+    parser.add_argument('--detect_anomaly', action='store_true', default=False)
+    parser.add_argument("--test_iterations", nargs="+", type=int, default=[1, 500, 800, 1000, 1500, 2000, 3000, 4000, 5000, 6000, 7_000, 30_000])
+    parser.add_argument("--save_iterations", nargs="+", type=int, default=[])
+    parser.add_argument("--quiet", action="store_true")
+    parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[])
+    parser.add_argument("--start_checkpoint", type=str, default = None)
+    parser.add_argument("--get_video", action="store_true")
+    parser.add_argument("--optim_pose", type=bool, default = True)
+    parser.add_argument("--gui", action="store_true")
+    parser.add_argument("--eval_pose", action="store_true")
+    parser.add_argument('--pose_eval_interval', type=int, default=100)
+    parser.add_argument('--psnr_threshold', type=float, default=26)
+    parser.add_argument('--gt_dynamic_mask', type=str, default='/home/remote/data/sintel/training/dynamic_label_perfect')
+    parser.add_argument('--dataset', type=str, default='sintel')
+
     
+    args = parser.parse_args(sys.argv[1:])
+    args.source_path = s
+    args.model_path = m
+    args.iterations = iter
+
+    args.save_iterations.append(args.iterations)
+
+    os.makedirs(args.model_path, exist_ok=True)
+    
+    print("Optimizing " + args.model_path)
+
+    # Initialize system state (RNG)
+    # safe_state(args.quiet)
+
+    torch.autograd.set_detect_anomaly(args.detect_anomaly)
+
+    if args.gui:
+        w, h = Image.open(os.path.join(args.source_path, 'images', 'frame_0000.png')).size
+        gui = GUI(gui = args.gui, w=w, h=h)
+    else:
+        gui = None
+
+    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, gui)
+
+    # All done
+    print("\nTraining complete.")
+                        
 if __name__ == "__main__":
     # Set up command line argument parser
     parser = ArgumentParser(description="Training script parameters")

{utils → utils_das3r}/camera_utils.py

@@ -11,8 +11,8 @@
 
 from scene.cameras import Camera
 import numpy as np
-from utils.general_utils import PILtoTorch
-from utils.graphics_utils import fov2focal
+from utils_das3r.general_utils import PILtoTorch
+from utils_das3r.graphics_utils import fov2focal
 import torch
 import scipy
 import matplotlib.pyplot as plt

{utils → utils_das3r}/pose_utils.py

@@ -3,7 +3,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 from typing import Tuple
-from utils.stepfun import sample_np, sample
+from utils_das3r.stepfun import sample_np, sample
 import scipy
 
 

{utils → utils_das3r}/rearrange.py

@@ -2,7 +2,7 @@ import os
 import cv2
 import PIL.Image as Image
 import numpy as np
-from vo_eval import file_interface
+from utils_das3r.vo_eval import file_interface
 from pathlib import Path
 from plyfile import PlyData, PlyElement
 import torch
@@ -38,7 +38,7 @@ def quaternion_to_matrix(quaternions):
     )
     return o.reshape(quaternions.shape[:-1] + (3, 3))
 
-def main(dataset_path, output_path):
+def rearrange(dataset_path, output_path):
     # load pred_traj
 
     output_colmap_path=os.path.join(output_path, 'sparse/0')
@@ -353,12 +353,13 @@ def R_to_quaternion(R):
 
     return np.array([w, x, y, z])
 
+def main(output_dir, rearranged_path):
+    rearrange(output_dir, rearranged_path)
+
 if __name__ == "__main__":
-    dataset_path = 'results/sintel'
-    output_path = dataset_path.replace('sintel', 'sintel_rearranged')
     parser = argparse.ArgumentParser(description='Rearrange dataset.')
     parser.add_argument('--output_dir', type=str, default='data/custom/output', help='Output directory')
     args = parser.parse_args()
     output_path = args.output_dir
     rearranged_path = output_path+'_rearranged'
-    main(args.output_dir, rearranged_path)
+    rearrange(args.output_dir, rearranged_path)