run_system.py

#
# Copyright (C) 2025, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import random
import sys
from queue import Empty
import shutil
import pypose as pp
import torch.nn.functional as F
from VSLAM.ImageFrame import Mode
from VSLAM.SharedKeyframes import SharedKeyframes
from VSLAM.SharedStates import SharedStates
from Reconstruct.gaussianviewer import GaussianViewer
from VSLAM.mvslam.visualization import run_visualization
from dataloaders.utils_load import load_dataset
from VSLAM.utils_config import load_config
from VSLAM.Frontend import Frontend
from plyfile import PlyElement, PlyData
from hydra.utils import instantiate
from safetensors.torch import load_file
from omegaconf import OmegaConf
import importlib.util

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import time

import numpy as np
import torch
from tqdm import tqdm
from socketserver import TCPServer
from http.server import SimpleHTTPRequestHandler
from Reconstruct.gaussianviewer import GaussianViewer
from Reconstruct.webviewer.webviewer import WebViewer
from threading import Thread
from graphdecoviewer.types import ViewerMode
import torch.multiprocessing as mp
import os

def load_args_from_path(file_path):
    # 1. 获取绝对路径并提取模块名
    abs_path = os.path.abspath(file_path)
    module_name = os.path.splitext(os.path.basename(file_path))[0]
    print(module_name, abs_path)

    # 2. 创建模块规范 (Spec)
    spec = importlib.util.spec_from_file_location(module_name, abs_path)
    
    # 3. 根据规范创建一个新的模块对象
    config_module = importlib.util.module_from_spec(spec)
    
    # 4. 执行模块以填充其内容
    spec.loader.exec_module(config_module)
    
    return config_module
def _inverse_sigmoid(tensor: torch.Tensor) -> torch.Tensor:
    return torch.log(tensor / (1.0 - tensor))

def convert_rgb_to_spherical_harmonics(rgb: torch.Tensor) -> torch.Tensor:
    """Convert RGB to degree-0 spherical harmonics.

    Reference:
        https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
    """
    coeff_degree0 = np.sqrt(1.0 / (4.0 * np.pi))
    return (rgb - 0.5) / coeff_degree0

def sixD2mtx(r):
    b1 = r[..., 0]
    b1 = b1 / torch.norm(b1, dim=-1, keepdim=True)
    b2 = r[..., 1] - torch.sum(b1 * r[..., 1], dim=-1, keepdim=True) * b1
    b2 = b2 / torch.norm(b2, dim=-1, keepdim=True)
    b3 = torch.cross(b1, b2, dim=-1)
    return torch.stack([b1, b2, b3], dim=-1)

if __name__ == "__main__":
    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    mp.set_start_method("spawn")
    torch.backends.cuda.matmul.allow_tf32 = False

    experiment_seed = int(os.environ.get("ARTDECO_EXPERIMENT_SEED", "0"))
    random.seed(experiment_seed)
    torch.manual_seed(experiment_seed)
    torch.cuda.manual_seed_all(experiment_seed)
    np.random.seed(experiment_seed)
    if len(sys.argv) < 2:
        print("Usage: python run_model.py <path_to_args_file>")
    else:
        base_model = sys.argv[1]
        args = getattr(
            __import__('config.args.'+base_model, fromlist=[''])
            , "get_args")()
    config = load_config(args.config)
    device_frontend = args.device_frontend
    device_backend = args.device_backend
    device_mapper = args.device_mapper
    device_shared = args.device_shared

    dataset = load_dataset(args)
    H_map, W_map = dataset.H_map, dataset.W_map
    H_slam, W_slam = dataset.H_slam, dataset.W_slam
    lc_count = {'big': 0, 'small': 0}

    # 2. Shared variables
    manager = mp.Manager()
    keyframes = SharedKeyframes(config, manager, H_slam, W_slam, buffer=max(len(dataset)+8, 32), device=device_shared)  # scene-sized storage
    states = SharedStates(manager, H_slam, W_slam, device=device_shared)  # shared

    model = instantiate(OmegaConf.load(args.model_conf), _recursive_=False).to(device_frontend)
    if args.model_weight.endswith(".pt"):
        model.load_state_dict(torch.load(args.model_weight, map_location='cpu', weights_only=True)['model'], strict=True)
    elif args.model_weight.endswith(".safetensors"):
        model.load_state_dict(load_file(args.model_weight, device=device_frontend), strict=True)
    else:
        raise ValueError('Expected a matching checkpoint ending in .pt or .safetensors')
        

    # 3. Initialize slam system
    print(f"Using dataset: {args.source_path}")
    frontend = Frontend(args, config, dataset, keyframes, states, model, device=device_frontend)
    frontend = mp.Process(target=frontend.run)
    frontend.start()

    # 4. Initialize the viewer
    
    scene_model = getattr(
        __import__('Reconstruct.scene.'+base_model+'.scene_models', fromlist=[''])
        , "SceneModel")(
            W_map, H_map, args, device=device_mapper
        )
    if args.viewer_mode in ["server", "local"]:
        viewer_mode = ViewerMode.SERVER if args.viewer_mode == "server" else ViewerMode.LOCAL
        viewer = GaussianViewer.from_scene_model(scene_model, viewer_mode, device = device_mapper)
        viewer_thd = Thread(target=viewer.run, args=(args.ip, args.port), daemon=True)
        viewer_thd.start()
        viewer.throttling = True # Enable throttling when training
    elif args.viewer_mode == "web":
        ip = "0.0.0.0"
        server = TCPServer((ip, 8000), SimpleHTTPRequestHandler)
        server_thd = Thread(target=server.serve_forever, daemon=True)
        server_thd.start()
        print(f"Visit http://{ip}:8000/webviewer to for the viewer")
        viewer = WebViewer(scene_model, args.ip, args.port)
        viewer_thd = Thread(target=viewer.run, daemon=True)
        viewer_thd.start()
    else:
        viewer = None

    # 5. Track and Reconstruct
    # Dict of runtimes for each step
    runtimes = ["Load", "BAB", "tri", "BAI", "Add", "Init", "Opt", "anc"]
    runtimes = {key: [0, 0] for key in runtimes}
    metrics = {}
    pbar = tqdm(total=0, desc=f"GS Mapper: ", position=1, leave=False)
    reconstruction_start_time = time.time()
    mapper_index = 0
    important_idx = 0
    related_frames = {}
    modules = __import__('Reconstruct.scene.'+base_model+'.keyframe', fromlist=[''])
    while True:
        mode = states.get_mode()
        try:
            keyframe_map_dict = states.msgGet()
        except Empty:
            if frontend.exitcode is not None and frontend.exitcode != 0:
                raise RuntimeError(f"Frontend failed with exit code {frontend.exitcode}")
            time.sleep(0.001)
            if mode == Mode.TERMINATED:
                break
            continue
        
        # construct keyframe
        frame_id = keyframe_map_dict["frame_id"]
        last_keyframe_index = keyframe_map_dict["last_keyframe_index"]
        last_keyframe_frame_id = keyframe_map_dict.get("last_keyframe_frame_id", None)
        if last_keyframe_index not in related_frames:
            related_frames[last_keyframe_index] = []
        related_frames[last_keyframe_index].append(mapper_index)
        
        T_CkC = keyframe_map_dict["T_CkC"]
        T_CkC = T_CkC.to(device_mapper) if T_CkC is not None else None
        original_img, info = dataset[frame_id]
        T_WC = pp.Sim3(keyframe_map_dict["T_WC"]).to(device_mapper)
        densePoint = keyframe_map_dict["densePoint"].to(device_mapper)
        K_slam = keyframe_map_dict["K"].to(device_mapper)
        is_slam_keyframe = keyframe_map_dict["is_slam_keyframe"]
        is_important = keyframe_map_dict["is_important"]
        is_test = keyframe_map_dict["is_test"]
        point_map = densePoint[...,:3].to(device_mapper) #  H W 3
        point_map_conf = densePoint[...,3].to(device_mapper) # H W
        Twc = pp.SE3(T_WC.data[:, :7])
        Twc = pp.quat2unit(Twc)
        Tcw = Twc.Inv().matrix()[0]

        image_for_mapper = dataset.transform.to_map(original_img, device=device_mapper)
        K_map = dataset.transform.K_slam2map(K_slam, device_mapper)

        keyframe_map = getattr(modules, 'Keyframe')(
            image_for_mapper,
            info["name"],
            is_test,
            Tcw,
            mapper_index,
            frame_id,
            last_keyframe_index,
            last_keyframe_frame_id,
            keyframe_map_dict["is_slam_keyframe"],
            K_map,
            args,
            T_CkCf=T_CkC,
            point_map=point_map.to(device_mapper),
            point_conf=point_map_conf.to(device_mapper),  # 1 1 H W
            device_mapper=device_mapper
        )

        with torch.cuda.device(device_mapper):
            scene_model.add_keyframe(keyframe_map)
            if is_important:
                scene_model.add_new_gaussians()
            num_iterations = args.num_key_iterations if is_important else args.num_common_iterations
            scene_model.optimization_loop(num_iterations, is_important)
            
            ## Intermediate evaluation
            if (
                len(scene_model.keyframes) % args.test_frequency == 0
                and args.test_frequency > 0
                and (args.test_hold > 0)
            ):
                metrics = scene_model.evaluate()

            ## Save intermediate model
            if (
                mapper_index % args.save_every == 0
                and args.save_every > 0
            ):
                scene_model.save(
                    os.path.join(args.model_path, "progress", f"{mapper_index:05d}")
                )

            ## Display optimization progress and metrics
            bar_postfix = []
            for key, value in metrics.items():
                bar_postfix += [f"\033[31m{key}:{value:.2f}\033[0m"]
            if args.display_runtimes:
                for key, value in runtimes.items():
                    if value[1] > 0:
                        bar_postfix += [
                            f"\033[35m{key}:{1000 * value[0] / value[1]:.1f}\033[0m"
                        ]
            bar_postfix += [
                # f"\033[36mFocal:{focal:.1f}",
                f"\033[36mKeyframes:{len(scene_model.keyframes)}\033[0m",
                f"\033[36mGaussians:{scene_model.n_active_gaussians}\033[0m",
            ]
            pbar.set_postfix_str(",".join(bar_postfix), refresh=False)
        mapper_index += 1
        pbar.n = frame_id
        pbar.refresh()

    if len(scene_model.keyframes) <= 0:
        time.sleep(5)
        exit(0)
    reconstruction_time = time.time() - reconstruction_start_time

    # Set to inference mode so that the model can be rendered properly
    scene_model.enable_inference_mode()

    # Save the model and metrics
    with torch.cuda.device(device_mapper):
        print("Saving the reconstruction to:", args.model_path)
        metrics = scene_model.save(args.model_path, reconstruction_time, len(dataset))
        # Fine tuning after initial reconstruction
        if len(args.save_at_finetune_epoch) > 0 or len(args.save_at_finetune_iteration) > 0:
            if len(args.save_at_finetune_epoch) > 0:
                finetune_epochs = max(args.save_at_finetune_epoch)
            else:
                finetune_epochs = int(max(args.save_at_finetune_iteration) / mapper_index)
                args.save_at_finetune_epoch = [finetune_epochs]
            exact_step_remainder = (max(args.save_at_finetune_iteration) % mapper_index
                                    if args.save_at_finetune_iteration else 0)
            torch.cuda.empty_cache()
            scene_model.inference_mode = False
            pbar = tqdm(range(0, finetune_epochs), desc="Fine tuning")
            for epoch in pbar:
                # Run one epoch of fine-tuning
                epoch_start_time = time.time()
                scene_model.finetune_epoch()
                if epoch + 1 == finetune_epochs:
                    for _ in range(exact_step_remainder):
                        scene_model.optimization_step(finetuning=True)
                epoch_time = time.time() - epoch_start_time
                reconstruction_time += epoch_time
                # Save the model and metrics
                if epoch + 1 in args.save_at_finetune_epoch:
                    torch.cuda.empty_cache()
                    scene_model.inference_mode = True
                    metrics = scene_model.save(
                        os.path.join(args.model_path, str(epoch + 1)), reconstruction_time
                    )
                    bar_postfix = []
                    for key, value in metrics.items():
                        bar_postfix += [f"\033[31m{key}:{value:.2f}\033[0m"]
                    pbar.set_postfix_str(",".join(bar_postfix))
                    scene_model.inference_mode = False
                    torch.cuda.empty_cache()

            # Set to inference mode so that the model can be rendered properly
            scene_model.inference_mode = True

        if args.save_to_data_for_gsplat:
            colmap_dir = os.path.join(args.model_path, "colmap")
            target_dir = os.path.join(args.source_path, "artdeco_colmap")
            if os.path.exists(target_dir):
                shutil.rmtree(target_dir)
            print(colmap_dir, target_dir)
            shutil.copytree(colmap_dir, target_dir)
        print(
            ", ".join(
                f"{metric}: {value:.3f}"
                if isinstance(value, float)
                else f"{metric}: {value}"
                for metric, value in metrics.items()
            )
        )
    if args.viewer_mode != "none":
        if args.viewer_mode == "web":
            while True:
                time.sleep(1)
        else:
            viewer.throttling = False  # Disable throttling when done training
            # Loop to keep the viewer alive
            while viewer.running:
                time.sleep(1)
    frontend.join()
    torch.cuda.empty_cache()
    del model