train.py

#
# Copyright (C) 2023, 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 numpy as np

import subprocess
cmd = 'nvidia-smi -q -d Memory |grep -A4 GPU|grep Used'
result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE).stdout.decode().split('\n')
os.environ['CUDA_VISIBLE_DEVICES']=str(np.argmin([int(x.split()[2]) for x in result[:-1]]))

os.system('echo $CUDA_VISIBLE_DEVICES')


import torch
import torchvision
import json
import wandb
import time
from os import makedirs
import shutil
from pathlib import Path
from PIL import Image
import torchvision.transforms.functional as tf
import lpips
from random import randint
from utils.loss_utils import l1_loss, ssim
from gaussian_renderer import prefilter_voxel, render, network_gui
import sys
from scene import Scene, GaussianModel
from utils.general_utils import safe_state
import uuid
from tqdm import tqdm
from utils.image_utils import psnr
from argparse import ArgumentParser, Namespace
from arguments import ModelParams, PipelineParams, OptimizationParams
import torch.nn.functional as F
# torch.set_num_threads(32)
lpips_fn = lpips.LPIPS(net='vgg').to('cuda')

try:
    from torch.utils.tensorboard import SummaryWriter
    TENSORBOARD_FOUND = True
    print("found tf board")
except ImportError:
    TENSORBOARD_FOUND = False
    print("not found tf board")

def saveRuntimeCode(dst: str) -> None:
    additionalIgnorePatterns = ['.git', '.gitignore']
    ignorePatterns = set()
    ROOT = '.'
    with open(os.path.join(ROOT, '.gitignore')) as gitIgnoreFile:
        for line in gitIgnoreFile:
            if not line.startswith('#'):
                if line.endswith('\n'):
                    line = line[:-1]
                if line.endswith('/'):
                    line = line[:-1]
                ignorePatterns.add(line)
    ignorePatterns = list(ignorePatterns)
    for additionalPattern in additionalIgnorePatterns:
        ignorePatterns.append(additionalPattern)

    log_dir = Path(__file__).resolve().parent

    shutil.copytree(log_dir, dst, ignore=shutil.ignore_patterns(*ignorePatterns))
    
    print('Backup Finished!')


import math

def depth_tolerance(iteration, total_iters, max_tol=1.0, min_tol=0.0, mode='cosine'):
    """
    根据迭代数动态调整 depth 容忍度
    iteration: 当前迭代数
    total_iters: 总迭代数
    max_tol: 最大容忍度
    min_tol: 最小容忍度
    mode: 'linear' 或 'cosine'
    """
    progress = min(1.0, iteration / total_iters)

    if mode == 'linear':
        tol = min_tol + (max_tol - min_tol) * progress
    elif mode == 'cosine':
        # 先快后慢，符合 curriculum learning
        tol = min_tol + (max_tol - min_tol) * (1 - math.cos(math.pi * progress)) / 2
    else:
        raise ValueError("mode must be 'linear' or 'cosine'")

    return tol


@torch.no_grad()
def depth_mask(gs_depth, depth_m, iteration, total_iters,
               max_tol=1.0, min_tol=0.05, mode='cosine'):
    """
    输入:
      gs_depth: [N] 每个高斯到相机的深度
      depth_m: [H,W] 深度图
    输出:
      mask: [N] 布尔，是否保留
    """
    tol = depth_tolerance(iteration, total_iters,
                          max_tol=max_tol, min_tol=min_tol, mode=mode)
    # 这里需要你有办法把 gs 对应到像素 (u,v)，拿到 depth_m[u,v]
    # 假设你已经有 matched_depth: [N] 每个GS对应像素的深度
    matched_depth = sample_depth_for_gs(gs_depth, depth_m)  # TODO: 你已有实现

    mask = (gs_depth >= matched_depth - tol) & (gs_depth <= matched_depth + tol)
    return mask


def resolve_depth_npy_dir(depth_npy_dir, model_path):
    if depth_npy_dir is None:
        depth_npy_dir = os.path.join(model_path, "mesh_depth_npy")
    return os.path.abspath(depth_npy_dir)


def load_depth_cache(cameras, depth_npy_dir):
    depth_cache = {}
    for camera in cameras:
        if camera.image_name in depth_cache:
            continue

        depth_path = os.path.join(depth_npy_dir, f"{camera.image_name}.npy")
        if not os.path.exists(depth_path):
            raise FileNotFoundError(f"Depth file not found for {camera.image_name}: {depth_path}")

        depth_cache[camera.image_name] = torch.from_numpy(np.load(depth_path)).float().squeeze()

    return depth_cache


def get_depth_for_camera(viewpoint_camera, depth_cache, device="cuda"):
    depth_m = depth_cache[viewpoint_camera.image_name]
    if depth_m.device.type == device:
        return depth_m
    return depth_m.to(device=device, non_blocking=True)


def training(dataset, opt, pipe, dataset_name, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from, wandb=None, logger=None, ply_path=None, mesh_path=None, depth_npy_dir=None):
    first_iter = 0
    tb_writer = prepare_output_and_logger(dataset)
    gaussians = GaussianModel(
        dataset.feat_dim, dataset.n_offsets, dataset.fork, dataset.use_feat_bank, dataset.appearance_dim, 
        dataset.add_opacity_dist, dataset.add_cov_dist, dataset.add_color_dist, dataset.add_level, 
        dataset.visible_threshold, dataset.dist2level, dataset.base_layer, dataset.progressive, dataset.extend
    )
    scene = Scene(dataset, gaussians, ply_path=ply_path, shuffle=False, logger=logger, resolution_scales=dataset.resolution_scales, mesh_path= mesh_path)
    gaussians.training_setup(opt)
    gaussians.set_coarse_interval(opt.coarse_iter, opt.coarse_factor)
    if checkpoint:
        (model_params, first_iter) = torch.load(checkpoint)
        gaussians.restore(model_params, opt)

    iter_start = torch.cuda.Event(enable_timing = True)
    iter_end = torch.cuda.Event(enable_timing = True)

    depth_npy_dir = resolve_depth_npy_dir(depth_npy_dir, dataset.model_path)
    all_cameras = scene.getTrainCameras() 
    depth_cache = load_depth_cache(all_cameras, depth_npy_dir)

    viewpoint_stack = None
    ema_loss_for_log = 0.0
    progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
    first_iter += 1
    # 7x7 平均卷积核
    kernel = torch.ones((1,1,7,7), device='cuda') / (7*7)
    epoch = 0
    
    for iteration in range(first_iter, opt.iterations + 1):        
        # network gui not available in octree-gs yet
        if network_gui.conn == None:
            network_gui.try_connect()
        while network_gui.conn != None:
            try:
                net_image_bytes = None
                custom_cam, do_training, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()
                if custom_cam != None:
                    net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"]
                    net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy())
                network_gui.send(net_image_bytes, dataset.source_path)
                if do_training and ((iteration < int(opt.iterations)) or not keep_alive):
                    break
            except Exception as e:
                network_gui.conn = None

        iter_start.record()

        gaussians.update_learning_rate(iteration)

        if dataset.random_background:
            bg_color = [np.random.random(),np.random.random(),np.random.random()] 
        elif dataset.white_background:
            bg_color = [1.0, 1.0, 1.0]
        else:
            bg_color = [0.0, 0.0, 0.0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        
        # Pick a random Camera
        if not viewpoint_stack:
            epoch = epoch + 1
            viewpoint_stack = scene.getTrainCameras().copy()
        viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))

        # Render
        if (iteration - 1) == debug_from:
            pipe.debug = True
        
        gaussians.set_anchor_mask(viewpoint_cam.camera_center, iteration, viewpoint_cam.resolution_scale)


        # max_tol = 1.5 
        # min_tol = 0

        # total_iters = 40_000

        # tol = depth_tolerance(iteration, total_iters,
        #                   max_tol=max_tol, min_tol=min_tol, mode='cosine')         



        depth_m = get_depth_for_camera(viewpoint_cam, depth_cache)
        voxel_visible_mask, depth_m = prefilter_voxel(viewpoint_cam, gaussians, pipe, background, depth_map=depth_m)


        ##add_dropout
        # true_indices = voxel_visible_mask.nonzero(as_tuple=True)[0]

        # if true_indices.numel() > 0:
        #     # 随机选择 2% 的 true 索引
        #     num_to_drop = max(1, int(0.02 * true_indices.numel()))
        #     sampled_idx = true_indices[torch.randperm(true_indices.numel())[:num_to_drop]]
        #     # 把这些位置设为 False
        #     voxel_visible_mask[sampled_idx] = False

        # false_indices = (~voxel_visible_mask).nonzero(as_tuple=True)[0]
        # if false_indices.numel() > 0:
        #     num_to_keep = max(1, int(0.05 * false_indices.numel()))  # 随机保留 5%
        #     sampled_idx = false_indices[torch.randperm(false_indices.numel())[:num_to_keep]]
        #     voxel_visible_mask[sampled_idx] = True
  




        retain_grad = (iteration < opt.update_until and iteration >= 0)
        render_pkg = render(viewpoint_cam, gaussians, pipe, background, visible_mask=voxel_visible_mask, retain_grad=retain_grad)
        
        image, viewspace_point_tensor, visibility_filter, offset_selection_mask, radii, scaling, opacity = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["selection_mask"], render_pkg["radii"], render_pkg["scaling"], render_pkg["neural_opacity"]

        # final_mask = voxel_visible_mask.clone()
        # final_mask[voxel_visible_mask] = mask_mesh
        # voxel_visible_mask = final_mask

        if viewpoint_cam.alpha_mask is not None:
            alpha_mask = viewpoint_cam.alpha_mask.cuda()
            image *= alpha_mask

        gt_image = viewpoint_cam.original_image.cuda()
        
        
        Ll1 = l1_loss(image, gt_image)

        
        
        # if  epoch  ==  11:
        #     if iteration % 1 == 0 :
        #         with torch.no_grad():
        #             pixel_loss_map = torch.abs((image - gt_image)).mean(dim=0).unsqueeze(0).unsqueeze(0)  # (1,1,H,W) 
                    
        #             patch_loss_map = F.conv2d(pixel_loss_map, kernel, stride=7, padding=3)  # (1,1,H,W)
                    
        #             patch_loss_map = patch_loss_map.squeeze(0).squeeze(0)   # (H//7, W//7)

        #             # 找到超过阈值的 patch 索引
        #             ys, xs = torch.where(patch_loss_map > 2*Ll1)

        #             # 映射回原图坐标 (patch 中心)
        #             patch_size = 7
        #             ys_pixel = ys * patch_size + patch_size // 2
        #             xs_pixel = xs * patch_size + patch_size // 2
    
            # gaussians.anchor_growing_by_mesh(xs_pixel, ys_pixel, depth_m, viewpoint_cam)

        ssim_loss = (1.0 - ssim(image, gt_image))
        if scaling.shape[0] > 0:
            scaling_reg = scaling.prod(dim=1).mean()
        # else:
        # scaling_reg = torch.tensor(0.0, device="cuda")
        loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * ssim_loss + 0.01*scaling_reg

        loss.backward()
        
        iter_end.record()
        anchor_number_use = voxel_visible_mask.sum()
        anchor_number_sum = gaussians._anchor.shape[0]

        with torch.no_grad():
            # Progress bar
            ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log

            if iteration % 10 == 0:
                progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})
                progress_bar.update(10)
            if iteration == opt.iterations:
                progress_bar.close()

            # Log and save
            training_report(tb_writer, dataset_name, iteration, anchor_number_use, anchor_number_sum, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background), wandb, logger, depth_cache)
            if (iteration in saving_iterations):
                logger.info("\n[ITER {}] Saving Gaussians".format(iteration))
                scene.save(iteration)
            
            # densification
            if iteration < opt.update_until and iteration > opt.start_stat:
                # add statis
                gaussians.training_statis(viewspace_point_tensor, opacity, visibility_filter, offset_selection_mask, voxel_visible_mask)
                
                # densification
                if opt.update_anchor and iteration > opt.update_from and iteration % opt.update_interval == 0:
                    # if epoch != 11 or 12 :
                    gaussians.adjust_anchor(
                        iteration=iteration,
                        check_interval=opt.update_interval, 
                        success_threshold=opt.success_threshold,
                        grad_threshold=opt.densify_grad_threshold, 
                        update_ratio=dataset.update_ratio,
                        extra_ratio=dataset.extra_ratio,
                        extra_up=dataset.extra_up,
                        min_opacity=opt.min_opacity 
                    )
            if iteration == 3 * len(scene.getTrainCameras()) or iteration == 6 * len(scene.getTrainCameras()):
                new_anchors_world = []
                for viewpoint_cam in scene.getTrainCameras():
                    gaussians.set_anchor_mask(viewpoint_cam.camera_center, iteration, viewpoint_cam.resolution_scale)
                    depth_m = get_depth_for_camera(viewpoint_cam, depth_cache)
                    voxel_visible_mask, depth_m = prefilter_voxel(viewpoint_cam, gaussians, pipe, background, depth_map=depth_m)
                    retain_grad = (iteration < opt.update_until and iteration >= 0)
                    render_pkg = render(viewpoint_cam, gaussians, pipe, background, visible_mask=voxel_visible_mask, retain_grad=retain_grad)
                    image = render_pkg["render"]
                    gt_image = viewpoint_cam.original_image.cuda()                    
                    Ll1 = l1_loss(image, gt_image)
                    
                    pixel_loss_map = torch.abs((image - gt_image)).mean(dim=0).unsqueeze(0).unsqueeze(0)  # (1,1,H,W) 
                    
                    patch_loss_map = F.conv2d(pixel_loss_map, kernel, stride=7, padding=3)  # (1,1,H,W)
                    
                    patch_loss_map = patch_loss_map.squeeze(0).squeeze(0)   # (H//7, W//7)

                    # 找到超过阈值的 patch 索引
                    ys, xs = torch.where(patch_loss_map > 3*Ll1)   ## 统计下loss相对于平均值的分布

                    # 映射回原图坐标 (patch 中心)
                    patch_size = 7
                    H, W = image.shape[-2:]
                    ys_pixel = (ys * patch_size + patch_size // 2).clamp(0, H-1)
                    xs_pixel = (xs * patch_size + patch_size // 2).clamp(0, W-1)
                    
                    device = gaussians._anchor.device

                    fx, fy, cx, cy = viewpoint_cam.Fx, viewpoint_cam.Fy, viewpoint_cam.Cx, viewpoint_cam.Cy
                    R = viewpoint_cam.R
                    t = viewpoint_cam.T

                    # === Step1: back-project (像素->相机系) ===
                    zs = depth_m[ys_pixel, xs_pixel]   # (N,)
                    xs = (xs_pixel - cx) * zs / fx
                    ys = (ys_pixel - cy) * zs / fy
                    pts_cam = torch.stack([xs, ys, zs], dim=1).to(device).float()

                    mask = torch.isfinite(pts_cam).all(dim=1) & torch.isfinite(zs) & (zs > 0)
                     
                    pts_cam_new = pts_cam[mask]

                    # === Step2: 相机->世界 (用 C2W) ===
                    Rt = np.zeros((4, 4), dtype=np.float32)
                    Rt[:3, :3] = R.transpose()
                    Rt[:3, 3] = t
                    Rt[3, 3] = 1.0

                    C2W = np.linalg.inv(Rt)   # camera->world
                    C2W = torch.from_numpy(C2W).to(device).float()

                    ones = torch.ones((pts_cam_new.shape[0], 1), device=device)
                    pts_cam_h = torch.cat([pts_cam_new, ones], dim=1)   # (M,4)
                    pts_world_h = (C2W @ pts_cam_h.T).T             # (M,4)
                    pts_world = pts_world_h[:, :3]                  # (M,3)        
                    new_anchors_world.append(pts_world)
                new_anchors_world = torch.cat(new_anchors_world, dim=0)
                gaussians.anchor_growing_by_mesh(new_anchors_world)
                    
                    
                        
            elif iteration == opt.update_until:
                del gaussians.opacity_accum
                del gaussians.offset_gradient_accum
                del gaussians.offset_denom
                torch.cuda.empty_cache()
            

                    
            # Optimizer step
            if iteration < opt.iterations:
                gaussians.optimizer.step()
                gaussians.optimizer.zero_grad(set_to_none = True)
            if (iteration in checkpoint_iterations):
                logger.info("\n[ITER {}] Saving Checkpoint".format(iteration))
                torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")

def prepare_output_and_logger(args):    
    if not args.model_path:
        if os.getenv('OAR_JOB_ID'):
            unique_str=os.getenv('OAR_JOB_ID')
        else:
            unique_str = str(uuid.uuid4())
        args.model_path = os.path.join("./output/", unique_str[0:10])
        
    # Set up output folder
    print("Output folder: {}".format(args.model_path))
    os.makedirs(args.model_path, exist_ok = True)
    with open(os.path.join(args.model_path, "cfg_args"), 'w') as cfg_log_f:
        cfg_log_f.write(str(Namespace(**vars(args))))

    # Create Tensorboard writer
    tb_writer = None
    if TENSORBOARD_FOUND:
        tb_writer = SummaryWriter(args.model_path)
    else:
        print("Tensorboard not available: not logging progress")
    return tb_writer

def training_report(tb_writer, dataset_name, iteration, anchor_number_use , anchor_number_sum, Ll1, loss, l1_loss, elapsed, testing_iterations, scene : Scene, renderFunc, renderArgs, wandb=None, logger=None, depth_cache=None):
    if tb_writer:
        tb_writer.add_scalar(f'{dataset_name}/train_loss_patches/l1_loss', Ll1.item(), iteration)
        tb_writer.add_scalar(f'{dataset_name}/train_loss_patches/total_loss', loss.item(), iteration)
        tb_writer.add_scalar(f'{dataset_name}/iter_time', elapsed, iteration)


    if wandb is not None:
        wandb.log({"train_l1_loss":Ll1, 'train_total_loss':loss, })
        wandb.log({"anchor_number_use":anchor_number_use, })
        wandb.log({"anchor_number_sum": anchor_number_sum, })
    
    # Report test and samples of training set
    if iteration in testing_iterations:
        scene.gaussians.eval()
        torch.cuda.empty_cache()
        
        validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()}, 
                                  {'name': 'train', 'cameras' : [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in range(5, 30, 5)]})

        for config in validation_configs:
            if config['cameras'] and len(config['cameras']) > 0:
                l1_test = 0.0
                psnr_test = 0.0
                
                # 创建保存图片的目录
                model_path = scene.model_path
                render_path = os.path.join(model_path, config['name'], "ours_{}".format(iteration), "renders")
                error_path = os.path.join(model_path, config['name'], "ours_{}".format(iteration), "errors")
                gts_path = os.path.join(model_path, config['name'], "ours_{}".format(iteration), "gt")
                makedirs(render_path, exist_ok=True)
                makedirs(error_path, exist_ok=True)
                makedirs(gts_path, exist_ok=True)
                
                if wandb is not None:
                    gt_image_list = []
                    render_image_list = []
                    errormap_list = []
                    

                for idx, viewpoint in enumerate(config['cameras']):
                    scene.gaussians.set_anchor_mask(viewpoint.camera_center, iteration, viewpoint.resolution_scale)
                    depth_m = get_depth_for_camera(viewpoint, depth_cache)
                    voxel_visible_mask, depth_m = prefilter_voxel(viewpoint, scene.gaussians, *renderArgs, depth_map=depth_m)
                    image = torch.clamp(renderFunc(viewpoint, scene.gaussians, *renderArgs, visible_mask=voxel_visible_mask)["render"], 0.0, 1.0)
                    gt_image = torch.clamp(viewpoint.original_image.to("cuda"), 0.0, 1.0)
                    
                    # 保存图片
                    if hasattr(viewpoint, 'image_name') and viewpoint.image_name:
                        # 使用 image_name 作为文件名（去除扩展名）
                        image_name = os.path.splitext(viewpoint.image_name)[0]
                        render_filename = os.path.join(render_path, image_name + ".png")
                        error_filename = os.path.join(error_path, image_name + ".png")
                        gt_filename = os.path.join(gts_path, image_name + ".png")
                    else:
                        # 使用索引作为文件名
                        render_filename = os.path.join(render_path, '{0:05d}'.format(idx) + ".png")
                        error_filename = os.path.join(error_path, '{0:05d}'.format(idx) + ".png")
                        gt_filename = os.path.join(gts_path, '{0:05d}'.format(idx) + ".png")
                    
                    # 确保 gt_image 格式正确（3通道）
                    if gt_image.dim() == 3:
                        gt_image_save = gt_image
                    else:
                        gt_image_save = gt_image[0:3, :, :] if gt_image.shape[0] > 3 else gt_image
                    
                    # 保存渲染图片、错误图和GT图片
                    torchvision.utils.save_image(image, render_filename)
                    torchvision.utils.save_image((gt_image_save - image).abs(), error_filename)
                    torchvision.utils.save_image(gt_image_save, gt_filename)
                    
                    if tb_writer and (idx < 30):
                        tb_writer.add_images(f'{dataset_name}/'+config['name'] + "_view_{}/render".format(viewpoint.image_name), image[None], global_step=iteration)
                        tb_writer.add_images(f'{dataset_name}/'+config['name'] + "_view_{}/errormap".format(viewpoint.image_name), (gt_image[None]-image[None]).abs(), global_step=iteration)

                        if wandb:
                            render_image_list.append(image[None])
                            errormap_list.append((gt_image[None]-image[None]).abs())
                            
                        if iteration == testing_iterations[0]:
                            tb_writer.add_images(f'{dataset_name}/'+config['name'] + "_view_{}/ground_truth".format(viewpoint.image_name), gt_image[None], global_step=iteration)
                            if wandb:
                                gt_image_list.append(gt_image[None])

                    l1_test += l1_loss(image, gt_image).mean().double()
                    psnr_test += psnr(image, gt_image).mean().double()

                
                
                psnr_test /= len(config['cameras'])
                l1_test /= len(config['cameras'])          
                logger.info("\n[ITER {}] Evaluating {}: L1 {} PSNR {}".format(iteration, config['name'], l1_test, psnr_test))

                
                if tb_writer:
                    tb_writer.add_scalar(f'{dataset_name}/'+config['name'] + '/loss_viewpoint - l1_loss', l1_test, iteration)
                    tb_writer.add_scalar(f'{dataset_name}/'+config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration)
                if wandb is not None:
                    wandb.log({f"{config['name']}_loss_viewpoint_l1_loss":l1_test, f"{config['name']}_PSNR":psnr_test})

        if tb_writer:
            # tb_writer.add_histogram(f'{dataset_name}/'+"scene/opacity_histogram", scene.gaussians.get_opacity, iteration)
            tb_writer.add_scalar(f'{dataset_name}/'+'total_points', scene.gaussians.get_anchor.shape[0], iteration)
        torch.cuda.empty_cache()

        scene.gaussians.train()

def render_set(model_path, name, iteration, views, gaussians, pipeline, background, depth_cache=None):
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    error_path = os.path.join(model_path, name, "ours_{}".format(iteration), "errors")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
    makedirs(render_path, exist_ok=True)
    makedirs(error_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)
    
    t_list = []
    visible_count_list = []
    per_view_dict = {}
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        
        torch.cuda.synchronize();t_start = time.time()
        
        gaussians.set_anchor_mask(view.camera_center, iteration, view.resolution_scale)
        depth_m = get_depth_for_camera(view, depth_cache)
        voxel_visible_mask, depth_m = prefilter_voxel(view, gaussians, pipeline, background, depth_map=depth_m)
        render_pkg = render(view, gaussians, pipeline, background, visible_mask=voxel_visible_mask)
        torch.cuda.synchronize();t_end = time.time()

        t_list.append(t_end - t_start)

        # renders
        rendering = torch.clamp(render_pkg["render"], 0.0, 1.0)
        visible_count = render_pkg["visibility_filter"].sum()
        visible_count_list.append(visible_count)

        # gts
        gt = view.original_image[0:3, :, :]
        
        # error maps
        if gt.device != rendering.device:
            rendering = rendering.to(gt.device)
        errormap = (rendering - gt).abs()

        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(errormap, os.path.join(error_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
        per_view_dict['{0:05d}'.format(idx) + ".png"] = visible_count.item()
        
    with open(os.path.join(model_path, name, "ours_{}".format(iteration), "per_view_count.json"), 'w') as fp:
            json.dump(per_view_dict, fp, indent=True)
    
    return t_list, visible_count_list

def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train=False, skip_test=False, wandb=None, tb_writer=None, dataset_name=None, logger=None, depth_npy_dir=None):
    with torch.no_grad():
        gaussians = GaussianModel(
            dataset.feat_dim, dataset.n_offsets, dataset.fork, dataset.use_feat_bank, dataset.appearance_dim, 
            dataset.add_opacity_dist, dataset.add_cov_dist, dataset.add_color_dist, dataset.add_level, 
            dataset.visible_threshold, dataset.dist2level, dataset.base_layer, dataset.progressive, dataset.extend
        )
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False, resolution_scales=dataset.resolution_scales)
        gaussians.eval()

        if dataset.random_background:
            bg_color = [np.random.random(),np.random.random(),np.random.random()] 
        elif dataset.white_background:
            bg_color = [1.0, 1.0, 1.0]
        else:
            bg_color = [0.0, 0.0, 0.0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        if not os.path.exists(dataset.model_path):
            os.makedirs(dataset.model_path)

        depth_npy_dir = resolve_depth_npy_dir(depth_npy_dir, dataset.model_path)
        all_cameras = scene.getTrainCameras() + scene.getTestCameras()
        depth_cache = load_depth_cache(all_cameras, depth_npy_dir)

        if not skip_train:
            t_train_list, visible_count  = render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background, depth_cache)
            train_fps = 1.0 / torch.tensor(t_train_list[5:]).mean()
            logger.info(f'Train FPS: \033[1;35m{train_fps.item():.5f}\033[0m')
            if wandb is not None:
                wandb.log({"train_fps":train_fps.item(), })

        if not skip_test:
            t_test_list, visible_count = render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background, depth_cache)
            test_fps = 1.0 / torch.tensor(t_test_list[5:]).mean()
            logger.info(f'Test FPS: \033[1;35m{test_fps.item():.5f}\033[0m')
            if tb_writer:
                tb_writer.add_scalar(f'{dataset_name}/test_FPS', test_fps.item(), 0)
            if wandb is not None:
                wandb.log({"test_fps":test_fps, })
    
    return visible_count


def readImages(renders_dir, gt_dir):
    renders = []
    gts = []
    image_names = []
    for fname in os.listdir(renders_dir):
        render = Image.open(renders_dir / fname)
        gt = Image.open(gt_dir / fname)
        renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda())
        gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda())
        image_names.append(fname)
    return renders, gts, image_names


def evaluate(model_paths, eval_name, visible_count=None, wandb=None, tb_writer=None, dataset_name=None, logger=None):

    full_dict = {}
    per_view_dict = {}
    full_dict_polytopeonly = {}
    per_view_dict_polytopeonly = {}
    print("")
    
    scene_dir = model_paths
    full_dict[scene_dir] = {}
    per_view_dict[scene_dir] = {}
    full_dict_polytopeonly[scene_dir] = {}
    per_view_dict_polytopeonly[scene_dir] = {}

    test_dir = Path(scene_dir) / eval_name

    for method in os.listdir(test_dir):

        full_dict[scene_dir][method] = {}
        per_view_dict[scene_dir][method] = {}
        full_dict_polytopeonly[scene_dir][method] = {}
        per_view_dict_polytopeonly[scene_dir][method] = {}

        method_dir = test_dir / method
        gt_dir = method_dir/ "gt"
        renders_dir = method_dir / "renders"
        renders, gts, image_names = readImages(renders_dir, gt_dir)

        ssims = []
        psnrs = []
        lpipss = []

        for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"):
            ssims.append(ssim(renders[idx], gts[idx]))
            psnrs.append(psnr(renders[idx], gts[idx]))
            lpipss.append(lpips_fn(renders[idx], gts[idx]).detach())
        
        if wandb is not None:
            wandb.log({"test_SSIMS":torch.stack(ssims).mean().item(), })
            wandb.log({"test_PSNR_final":torch.stack(psnrs).mean().item(), })
            wandb.log({"test_LPIPS":torch.stack(lpipss).mean().item(), })

        logger.info(f"model_paths: \033[1;35m{model_paths}\033[0m")
        logger.info("  SSIM : \033[1;35m{:>12.7f}\033[0m".format(torch.tensor(ssims).mean(), ".5"))
        logger.info("  PSNR : \033[1;35m{:>12.7f}\033[0m".format(torch.tensor(psnrs).mean(), ".5"))
        logger.info("  LPIPS: \033[1;35m{:>12.7f}\033[0m".format(torch.tensor(lpipss).mean(), ".5"))
        print("")


        if tb_writer:
            tb_writer.add_scalar(f'{dataset_name}/SSIM', torch.tensor(ssims).mean().item(), 0)
            tb_writer.add_scalar(f'{dataset_name}/PSNR', torch.tensor(psnrs).mean().item(), 0)
            tb_writer.add_scalar(f'{dataset_name}/LPIPS', torch.tensor(lpipss).mean().item(), 0)
            
            tb_writer.add_scalar(f'{dataset_name}/VISIBLE_NUMS', torch.tensor(visible_count).mean().item(), 0)
        
        full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(),
                                                "PSNR": torch.tensor(psnrs).mean().item(),
                                                "LPIPS": torch.tensor(lpipss).mean().item()})
        per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)},
                                                    "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)},
                                                    "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)},
                                                    "VISIBLE_COUNT": {name: vc for vc, name in zip(torch.tensor(visible_count).tolist(), image_names)}})

    with open(scene_dir + "/results.json", 'w') as fp:
        json.dump(full_dict[scene_dir], fp, indent=True)
    with open(scene_dir + "/per_view.json", 'w') as fp:
        json.dump(per_view_dict[scene_dir], fp, indent=True)
    
def get_logger(path):
    import logging

    logger = logging.getLogger()
    logger.setLevel(logging.INFO) 
    fileinfo = logging.FileHandler(os.path.join(path, "outputs.log"))
    fileinfo.setLevel(logging.INFO) 
    controlshow = logging.StreamHandler()
    controlshow.setLevel(logging.INFO)
    formatter = logging.Formatter("%(asctime)s - %(levelname)s: %(message)s")
    fileinfo.setFormatter(formatter)
    controlshow.setFormatter(formatter)

    logger.addHandler(fileinfo)
    logger.addHandler(controlshow)

    return logger

if __name__ == "__main__":
    # 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('--warmup', action='store_true', default=False)
    parser.add_argument('--use_wandb', action='store_true', default=False)
    parser.add_argument("--test_iterations", nargs="+", type=int, default=[-1])
    parser.add_argument("--save_iterations", nargs="+", type=int, default=[30_000, 50_000])
    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("--gpu", type=str, default = '-1')
    parser.add_argument("--ply_path", type=str, default=None)
    parser.add_argument("--ply_mesh", type=str, default=None)
    parser.add_argument("--depth_npy_dir", type=str, default=None)
    args = parser.parse_args(sys.argv[1:])

    # enable logging
    model_path = args.model_path
    os.makedirs(model_path, exist_ok=True)

    logger = get_logger(model_path)

    logger.info(f'args: {args}')

    if args.test_iterations[0] == -1:
        args.test_iterations = [i for i in range(10000, args.iterations + 1, 10000)]
    if len(args.test_iterations) == 0 or args.test_iterations[-1] != args.iterations:
        args.test_iterations.append(args.iterations)
    print(args.test_iterations)

    if args.save_iterations[0] == -1:
        args.save_iterations = [i for i in range(10000, args.iterations + 1, 10000)]
    if len(args.save_iterations) == 0 or args.save_iterations[-1] != args.iterations:
        args.save_iterations.append(args.iterations)
    print(args.save_iterations)

    if args.gpu != '-1':
        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
        os.system("echo $CUDA_VISIBLE_DEVICES")
        logger.info(f'using GPU {args.gpu}')

    # try:
    #     saveRuntimeCode(os.path.join(args.model_path, 'backup'))
    # except:
    #     logger.info(f'save code failed~')
        
    dataset = args.source_path.split('/')[-1]
    exp_name = args.model_path.split('/')[-1]
    
    if args.use_wandb:
        wandb.login(key='1a21dba66d9736777e51aa1700ab09d6623a9183')
        wandb.login(verify=False) 
        run = wandb.init(
            # Set the project where this run will be logged
            project=f"least-gs",
            name=exp_name,
            # Track hyperparameters and run metadata
            settings=wandb.Settings(start_method="fork"),
            config=vars(args)
        )
    else:
        wandb = None
    
    logger.info("Optimizing " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    # Start GUI server, configure and run training
    # network_gui.init(args.ip, args.port)
    torch.autograd.set_detect_anomaly(args.detect_anomaly)
    args.depth_npy_dir = resolve_depth_npy_dir(args.depth_npy_dir, args.model_path)
    
    # training
    training(lp.extract(args), op.extract(args), pp.extract(args), dataset,  args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from, wandb, logger, args.ply_path, mesh_path=args.ply_mesh, depth_npy_dir=args.depth_npy_dir)
    if args.warmup:
        logger.info("\n Warmup finished! Reboot from last checkpoints")
        new_ply_path = os.path.join(args.model_path, f'point_cloud/iteration_{args.iterations}', 'point_cloud.ply')
        training(lp.extract(args), op.extract(args), pp.extract(args), dataset,  args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from, wandb=wandb, logger=logger, ply_path=new_ply_path, mesh_path=args.ply_mesh, depth_npy_dir=args.depth_npy_dir)

    # All done
    logger.info("\nTraining complete.")

    # rendering
    logger.info(f'\nStarting Rendering~')
    if args.eval:
        visible_count = render_sets(lp.extract(args), -1, pp.extract(args), skip_train=True, skip_test=False, wandb=wandb, logger=logger, depth_npy_dir=args.depth_npy_dir)
    else:
        visible_count = render_sets(lp.extract(args), -1, pp.extract(args), skip_train=False, skip_test=True, wandb=wandb, logger=logger, depth_npy_dir=args.depth_npy_dir)
    logger.info("\nRendering complete.")

    # calc metrics
    logger.info("\n Starting evaluation...")
    eval_name = 'test' if args.eval else 'train'
    evaluate(args.model_path, eval_name, visible_count=visible_count, wandb=wandb, logger=logger)
    logger.info("\nEvaluating complete.")