run_upst.py

# import imp
import os, sys, copy, glob, json, time, random, argparse


os.environ["CUDA_VISIBLE_DEVICES"]="0"

from shutil import copyfile
from tqdm import tqdm, trange
from PIL import Image
from PIL import ImageFile
import mmcv
import imageio
import numpy as np
from torch.nn import functional as F
import torch
import torch.nn as nn
import torch.nn.functional as F

from lib import utils, upst, upstmpi
from lib.load_data import load_data
from torchvision import transforms
from torchvision.utils import save_image

from YUVStyleNet import *
def find_files(dir, exts):
    files_grabbed = []
    for ext in exts:
        files_grabbed.extend(glob.glob(os.path.join(dir, ext)))
    if len(files_grabbed) > 0:
        files_grabbed = sorted(files_grabbed)
    return files_grabbed

data_transform = transforms.Compose([
    transforms.RandomResizedCrop(256, scale=(256/480, 1), ratio=(1, 1)),
    transforms.ToTensor(),
])

data_transform_view = transforms.Compose([
    transforms.Resize(size=(256, 256)),
    transforms.ToTensor(),
])

def get_style_input(style_imgs, mode=None, test_seed=None, style_ID=None):
    if mode == "test":
        random.seed(test_seed)
        if style_ID != None:
            style_img_path = style_imgs[style_ID]
        else:
            style_img_path = random.sample(style_imgs, 1)[0]
    else:
        style_img_path = np.random.choice(style_imgs, 1)[0]
    
    ori_style_img = Image.open(style_img_path).convert('RGB')
    style_img = data_transform(ori_style_img)
    style_idx = torch.from_numpy(np.array([style_imgs.index(style_img_path)]))
    
    return style_img, style_idx

# Style loss and Content loss
def get_style_loss(style_feats, transformed_feats, style_weight=15):
    return F.mse_loss(transformed_feats, style_feats) * style_weight

def get_content_loss(content_feat, transformed_feat, content_weight=1):
    return F.mse_loss(transformed_feat, content_feat) * content_weight

def config_parser():
    '''Define command line arguments
    '''

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config', default="./configs/nerf/chair.py",
    # parser.add_argument('--config', default="./configs/llff/fern.py",
                        help='config file path')
    parser.add_argument("--style_img",default="./style_images/IMG_20210921_192826.jpg",
                        help='style_img file path')
    parser.add_argument("--seed", type=int, default=777,
                        help='Random seed')
    parser.add_argument("--no_reload", action='store_true',
                        help='do not reload weights from saved ckpt')
    parser.add_argument("--no_reload_optimizer", action='store_true',
                        help='do not reload optimizer state from saved ckpt')
    parser.add_argument("--ft_path", type=str, default='',
                        help='specific weights npy file to reload for coarse network')
    parser.add_argument("--export_bbox_and_cams_only", type=str, default='',
                        help='export scene bbox and camera poses for debugging and 3d visualization')
    parser.add_argument("--export_coarse_only", type=str, default='')

    # testing options
    parser.add_argument("--style_base_root",default="./data/coco/mscoco2017")
    parser.add_argument("--render_style", action='store_true',default=False)
    parser.add_argument("--render_geometry", action='store_true',default=True)
    parser.add_argument("--render_only", action='store_true',default=False,
                        help='do not optimize, reload weights and render out render_poses path')
    parser.add_argument("--render_test", action='store_true',default=False)
    parser.add_argument("--render_train", action='store_true',default=False)
    parser.add_argument("--render_video", action='store_true',default=True)
    parser.add_argument("--render_video_factor", type=int, default=0,
                        help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
    parser.add_argument("--eval_ssim", action='store_true',default=True)
    parser.add_argument("--eval_lpips_alex", action='store_true',default=False)
    parser.add_argument("--eval_lpips_vgg", action='store_true',default=False)

    # logging/saving options
    parser.add_argument("--i_print",   type=int, default=500,
                        help='frequency of console printout and metric loggin')
    parser.add_argument("--i_weights", type=int, default=10000,
                        help='frequency of weight ckpt saving')
    return parser


@torch.no_grad()
def render_viewpoints_style(style_img, model, render_poses, HW, Ks, ndc, render_kwargs,
                      gt_imgs=None, savedir=None, render_factor=0,
                      eval_ssim=False, eval_lpips_alex=False, eval_lpips_vgg=False):
    '''Render images for the given viewpoints; run evaluation if gt given.
    '''
    assert len(render_poses) == len(HW) and len(HW) == len(Ks)
    model.eval()
    print(style_img)
    style_images = Image.open(os.path.join(style_img)).convert('RGB')

    style_embedded = data_transform_view(style_images)
    style_embedded = style_embedded.to(device)
    if render_factor!=0:
        HW = np.copy(HW)
        Ks = np.copy(Ks)
        HW //= render_factor
        Ks[:, :2, :3] //= render_factor

    rgbs = []
    depths = []
    psnrs = []
    ssims = []
    lpips_alex = []
    lpips_vgg = []

    for i, c2w in enumerate(tqdm(render_poses)):

        H, W = HW[i]
        K = Ks[i]
        rays_o, rays_d, viewdirs = upst.get_rays_of_a_view(
                H, W, K, c2w, ndc, inverse_y=render_kwargs['inverse_y'],
                flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
        keys = ['rgb_marched', 'depth']
        rays_o = rays_o.flatten(0,-2)
        rays_d = rays_d.flatten(0,-2)
        viewdirs = viewdirs.flatten(0,-2)

        render_result_chunks = [
            {k: v for k, v in model(ro, rd, vd, torch.unsqueeze(style_embedded, 0), **render_kwargs).items() if k in keys}
            for ro, rd, vd in zip(rays_o.split(100*100, 0), rays_d.split(100*100, 0), viewdirs.split(100*100, 0))
        ]
        render_result = {
            k: torch.cat([ret[k] for ret in render_result_chunks]).reshape(H,W,-1)
            for k in render_result_chunks[0].keys()
        }


        rgb = render_result['rgb_marched'].cpu().numpy()
        depth = render_result['depth'].cpu().numpy()
        rgbs.append(rgb)
        depths.append(depth)
        if gt_imgs is not None and render_factor==0:
            p = -10. * np.log10(np.mean(np.square(rgb - gt_imgs[i])))
            psnrs.append(p)
            if eval_ssim:
                ssims.append(utils.rgb_ssim(rgb, gt_imgs[i], max_val=1))
            if eval_lpips_alex:
                lpips_alex.append(utils.rgb_lpips(rgb, gt_imgs[i], net_name='alex', device=c2w.device))
            if eval_lpips_vgg:
                lpips_vgg.append(utils.rgb_lpips(rgb, gt_imgs[i], net_name='vgg', device=c2w.device))

    if len(psnrs):
        print('Testing psnr', np.mean(psnrs), '(avg)')
        if eval_ssim: print('Testing ssim', np.mean(ssims), '(avg)')
        if eval_lpips_vgg: print('Testing lpips (vgg)', np.mean(lpips_vgg), '(avg)')
        if eval_lpips_alex: print('Testing lpips (alex)', np.mean(lpips_alex), '(avg)')

    if savedir is not None:
        print(f'Writing images to {savedir}')
        for i in trange(len(rgbs)):
            rgb8 = utils.to8b(rgbs[i])
            filename = os.path.join(savedir, '{:03d}.png'.format(i))
            imageio.imwrite(filename, rgb8)

    rgbs = np.array(rgbs)
    depths = np.array(depths)

    return rgbs, depths



@torch.no_grad()
def render_viewpoints(model, render_poses, HW, Ks, ndc, render_kwargs,
                      gt_imgs=None, savedir=None, render_factor=0,
                      eval_ssim=False, eval_lpips_alex=False, eval_lpips_vgg=False):
    '''Render images for the given viewpoints; run evaluation if gt given.
    '''
    assert len(render_poses) == len(HW) and len(HW) == len(Ks)
    model.eval()

    if render_factor!=0:
        HW = np.copy(HW)
        Ks = np.copy(Ks)
        HW //= render_factor
        Ks[:, :2, :3] //= render_factor

    rgbs = []
    depths = []
    psnrs = []
    ssims = []
    lpips_alex = []
    lpips_vgg = []

    for i, c2w in enumerate(tqdm(render_poses)):

        H, W = HW[i]
        K = Ks[i]
        rays_o, rays_d, viewdirs = upst.get_rays_of_a_view(
                H, W, K, c2w, ndc, inverse_y=render_kwargs['inverse_y'],
                flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
        keys = ['rgb_marched', 'depth']
        rays_o = rays_o.flatten(0,-2)
        rays_d = rays_d.flatten(0,-2)
        viewdirs = viewdirs.flatten(0,-2)
        # style_embedded = None

        render_result_chunks = [
            {k: v for k, v in model(ro, rd, vd, None, **render_kwargs).items() if k in keys}
            for ro, rd, vd in zip(rays_o.split(100*100, 0), rays_d.split(100*100, 0), viewdirs.split(100*100, 0))
        ]
        render_result = {
            k: torch.cat([ret[k] for ret in render_result_chunks]).reshape(H,W,-1)
            for k in render_result_chunks[0].keys()
        }


        rgb = render_result['rgb_marched'].cpu().numpy()
        depth = render_result['depth'].cpu().numpy()
        rgbs.append(rgb)
        depths.append(depth)
        if i==0:
            print('Testing', rgb.shape)

        if gt_imgs is not None and render_factor==0:
            p = -10. * np.log10(np.mean(np.square(rgb - gt_imgs[i])))
            psnrs.append(p)
            if eval_ssim:
                ssims.append(utils.rgb_ssim(rgb, gt_imgs[i], max_val=1))
            if eval_lpips_alex:
                lpips_alex.append(utils.rgb_lpips(rgb, gt_imgs[i], net_name='alex', device=c2w.device))
            if eval_lpips_vgg:
                lpips_vgg.append(utils.rgb_lpips(rgb, gt_imgs[i], net_name='vgg', device=c2w.device))

    if len(psnrs):
        print('Testing psnr', np.mean(psnrs), '(avg)')
        if eval_ssim: print('Testing ssim', np.mean(ssims), '(avg)')
        if eval_lpips_vgg: print('Testing lpips (vgg)', np.mean(lpips_vgg), '(avg)')
        if eval_lpips_alex: print('Testing lpips (alex)', np.mean(lpips_alex), '(avg)')

    if savedir is not None:
        print(f'Writing images to {savedir}')
        for i in trange(len(rgbs)):
            rgb8 = utils.to8b(rgbs[i])
            filename = os.path.join(savedir, '{:03d}.png'.format(i))
            imageio.imwrite(filename, rgb8)

    rgbs = np.array(rgbs)
    depths = np.array(depths)

    return rgbs, depths



def seed_everything():
    '''Seed everything for better reproducibility.
    (some pytorch operation is non-deterministic like the backprop of grid_samples)
    '''
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

def load_everything(cfg):
    '''Load images / poses / camera settings / data split.
    '''
    data_dict = load_data(cfg.data)

    # remove useless field
    kept_keys = {
            'hwf', 'HW', 'Ks', 'near', 'far',
            'i_train', 'i_val', 'i_test', 'irregular_shape',
            'poses', 'render_poses', 'images'}
    for k in list(data_dict.keys()):
        if k not in kept_keys:
            data_dict.pop(k)

    # construct data tensor
    if data_dict['irregular_shape']:
        data_dict['images'] = [torch.FloatTensor(im, device='cpu') for im in data_dict['images']]
    else:
        data_dict['images'] = torch.FloatTensor(data_dict['images'], device='cpu')
    data_dict['poses'] = torch.Tensor(data_dict['poses'])
    return data_dict


def compute_bbox_by_cam_frustrm(args, cfg, HW, Ks, poses, i_train, near, far, **kwargs):
    print('compute_bbox_by_cam_frustrm: start')
    xyz_min = torch.Tensor([np.inf, np.inf, np.inf])
    xyz_max = -xyz_min
    for (H, W), K, c2w in zip(HW[i_train], Ks[i_train], poses[i_train]):
        rays_o, rays_d, viewdirs = upst.get_rays_of_a_view(
                H=H, W=W, K=K, c2w=c2w,
                ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
        if cfg.data.ndc:
            pts_nf = torch.stack([rays_o+rays_d*near, rays_o+rays_d*far])
        else:
            pts_nf = torch.stack([rays_o+viewdirs*near, rays_o+viewdirs*far])
        xyz_min = torch.minimum(xyz_min, pts_nf.amin((0,1,2)))
        xyz_max = torch.maximum(xyz_max, pts_nf.amax((0,1,2)))
    print('compute_bbox_by_cam_frustrm: xyz_min', xyz_min)
    print('compute_bbox_by_cam_frustrm: xyz_max', xyz_max)
    print('compute_bbox_by_cam_frustrm: finish')
    return xyz_min, xyz_max

@torch.no_grad()
def compute_bbox_by_coarse_geo(model_class, model_path, thres):
    print('compute_bbox_by_coarse_geo: start')
    eps_time = time.time()
    model = utils.load_model(model_class, model_path)
    interp = torch.stack(torch.meshgrid(
        torch.linspace(0, 1, model.density.shape[2]),
        torch.linspace(0, 1, model.density.shape[3]),
        torch.linspace(0, 1, model.density.shape[4]),
    ), -1)
    dense_xyz = model.xyz_min * (1-interp) + model.xyz_max * interp
    density = model.grid_sampler(dense_xyz, model.density)
    alpha = model.activate_density(density)
    mask = (alpha > thres)
    active_xyz = dense_xyz[mask]
    xyz_min = active_xyz.amin(0)
    xyz_max = active_xyz.amax(0)
    print('compute_bbox_by_coarse_geo: xyz_min', xyz_min)
    print('compute_bbox_by_coarse_geo: xyz_max', xyz_max)
    eps_time = time.time() - eps_time
    print('compute_bbox_by_coarse_geo: finish (eps time:', eps_time, 'secs)')
    return xyz_min, xyz_max

class AverageMeter(object):
    """
    Keeps track of most recent, average, sum, and count of a metric.
    """

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def scene_rep_reconstruction_style(args, cfg, cfg_model, cfg_train, data_dict, stage, coarse_ckpt_path=None):

    style_base_root = args.style_base_root
    style_dir = os.path.join(style_base_root, "train2017")
    style_img_files = find_files(style_dir, exts=['*.png', '*.jpg'])

    style_dir = os.path.join(style_base_root, "test2017")
    style_img_files+= find_files(style_dir, exts=['*.png', '*.jpg'])

    # logger.info("Number of style images is {}".format(len(style_img_files)))
    print("Number of style images is {}".format(len(style_img_files)))

    # init
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    HW, Ks, near, far, i_train, i_val, i_test, poses, render_poses, images = [
        data_dict[k] for k in [
            'HW', 'Ks', 'near', 'far', 'i_train', 'i_val', 'i_test', 'poses', 'render_poses', 'images'
        ]
    ]

    last_ckpt_path = os.path.join(cfg.basedir, cfg.expname, f'{stage}_last.tar')
    if args.no_reload:
        reload_ckpt_path = None
    elif args.ft_path:
        reload_ckpt_path = args.ft_path
    elif os.path.isfile(last_ckpt_path):
        reload_ckpt_path = last_ckpt_path
    else:
        reload_ckpt_path = None

    # init model and optimizer
    if reload_ckpt_path is None:
        print(f'scene_rep_reconstruction ({stage}): reload from {coarse_ckpt_path}')
        if cfg.data.ndc:
            model_class = upstmpi.UPSTMPI_DirectGO
        else:
            model_class = upst.UPST_DirectVoxGO
        model = utils.load_model(model_class, coarse_ckpt_path).to(device)
        optimizer = utils.create_optimizer_or_freeze_model(model, cfg_train, global_step=0)
    else:
        print(f'scene_rep_reconstruction ({stage}): reload from {reload_ckpt_path}')
        if cfg.data.ndc:
            model_class = upstmpi.UPSTMPI_DirectGO
        else:
            model_class = upst.UPST_DirectVoxGO
        model = utils.load_model(model_class, reload_ckpt_path).to(device)
        optimizer = utils.create_optimizer_or_freeze_model(model, cfg_train, global_step=0)
        model, optimizer, start = utils.load_checkpoint(
                model, optimizer, reload_ckpt_path, args.no_reload_optimizer)

    # init rendering setup
    render_kwargs = {
        'near': data_dict['near'],
        'far': data_dict['far'],
        'bg': 1 if cfg.data.white_bkgd else 0,
        'stepsize': cfg_model.stepsize,
        'inverse_y': cfg.data.inverse_y,
        'flip_x': cfg.data.flip_x,
        'flip_y': cfg.data.flip_y,
    }

    rays_o_tr=[]
    rays_d_tr=[]
    viewdirs_tr=[]
    for (H, W), K, c2w in zip(HW[i_train], Ks[i_train], poses[i_train]):
        rays_o, rays_d, viewdirs = upst.get_rays_of_a_view(#得到所有的光线起点和角度,以及观测角度,(W,H,3)->([800, 800, 3])
                H, W, K, c2w, ndc=cfg.data.ndc, inverse_y=render_kwargs['inverse_y'],
                flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
        rays_o_tr.append(rays_o)
        rays_d_tr.append(rays_d)
        viewdirs_tr.append(viewdirs)


    if cfg_train.pervoxel_lr:
        def per_voxel_init():
            cnt = model.voxel_count_views(
                    rays_o_tr=rays_o_tr, rays_d_tr=rays_d_tr, imsz=imsz, near=near, far=far,
                    stepsize=cfg_model.stepsize, downrate=cfg_train.pervoxel_lr_downrate,
                    irregular_shape=data_dict['irregular_shape'])
            optimizer.set_pervoxel_lr(cnt)
            with torch.no_grad():
                model.density[cnt <= 2] = -100
        per_voxel_init()




    glow = YUVStyleNet(3)
    print(glow)
    print('Total params: %.2fM' % (sum(p.numel() for p in glow.parameters()) / 1000000.0))
    decoderdefault='./YUVStyleNet.pth'
    if os.path.isfile(decoderdefault):
        print("--------loading checkpoint----------")
        checkpoint = torch.load(decoderdefault)
        args.start_iter = checkpoint['iter']
        glow.load_state_dict(checkpoint['state_dict'])
        print("=> loaded checkpoint '{}'".format(decoderdefault))
    else:
        print("--------no checkpoint found---------")

    glow = glow.to(device)
    glow.eval()





    losses = AverageMeter()
    c_losses = AverageMeter()
    s_losses = AverageMeter()


    # GOGO
    torch.cuda.empty_cache()
    psnr_lst = []
    time0 = time.time()
    global_step = -1
    # start = 0
    for global_step in trange(1+start, 1+cfg_train.N_iters):
        random_i = np.random.randint(len(i_train))
        # random_i=0
        target_rgb = images[i_train][random_i]
        c2w = poses[random_i]

        rays_o = rays_o_tr[random_i].flatten(0,-2)
        rays_d = rays_d_tr[random_i].flatten(0,-2)
        viewdirs = viewdirs_tr[random_i].flatten(0,-2)

        total_inds = np.arange(H*W).reshape(H, W)
        patch_H, patch_W = 100, 100
        num_region_H, num_region_W = H//patch_H, W//patch_W
        
        region_size_v = np.random.randint(num_region_H//2, num_region_H+1)
        region_size_u = np.random.randint(num_region_W//3, num_region_W+1)
        region_position_v = np.random.randint(H - patch_H * region_size_v + region_size_v)
        region_position_u = np.random.randint(W - patch_W * region_size_u + region_size_u)
        select_inds = total_inds[region_position_v::region_size_v][:patch_H][:, region_position_u::region_size_u][:, :patch_W].reshape(-1)
        rays_o_temp = rays_o[select_inds, :]
        rays_d_temp = rays_d[select_inds, :]
        viewdirs_temp = viewdirs[select_inds, :]
        target_rgb = target_rgb.flatten(0,-2)[select_inds, :]  




        if global_step < 0:
            style_embedded = None
            optimizer.zero_grad()
            render_result = model(rays_o_temp, rays_d_temp, viewdirs_temp,style_embedded, global_step=global_step, **render_kwargs)
            # gradient descent step
            target = target_rgb.flatten(0,-2)
            target = target_rgb.to(device)

            rgb_loss = cfg_train.weight_main * utils.img2mse(render_result['rgb_marched'], target)
            psnr = utils.mse2psnr(rgb_loss.detach())
            loss = rgb_loss
            if cfg_train.weight_entropy_last > 0:
                pout = render_result['alphainv_last'].clamp(1e-6, 1-1e-6)
                entropy_last_loss = -(pout*torch.log(pout) + (1-pout)*torch.log(1-pout)).mean()
                loss += cfg_train.weight_entropy_last * entropy_last_loss
            if cfg_train.weight_rgbper > 0:
                rgbper = (render_result['raw_rgb'] - target[render_result['ray_id']]).pow(2).sum(-1)
                rgbper_loss = (rgbper * render_result['weights'].detach()).sum() / len(rays_o)
                loss += cfg_train.weight_rgbper * rgbper_loss            
        else:
            #1
            model.k0.requires_grad = False
            model.density.requires_grad = False
            style_img, style_idx= get_style_input(style_imgs=style_img_files, mode=None, test_seed=None, style_ID=None)
            style_embedded = style_img.cuda()
            target_rgb = target_rgb.to(device)
            # optimizer.zero_grad(set_to_none=True)
            optimizer.zero_grad()
            render_result = model(rays_o_temp, rays_d_temp, viewdirs_temp,torch.unsqueeze(style_embedded, 0), global_step=global_step, **render_kwargs)





            # # gradient descent step
            alphainv_last = render_result['alphainv_last'].reshape(patch_H,patch_W,1)
            alphainv_last = alphainv_last.repeat_interleave(3, dim=2)


            zero_temp = torch.zeros_like(alphainv_last)
            alphainv_last = torch.where(alphainv_last < 1.0, zero_temp, alphainv_last)


            style_img = torch.unsqueeze(style_img, 0)
            style_img = style_img.to(device)
            content_images = torch.unsqueeze(target_rgb.reshape(patch_H,patch_W,3).permute(2,0,1), 0)
            content_images = F.interpolate(content_images, (style_img.size()[2], style_img.size()[3]),  mode='bilinear', align_corners=False)


            s_output, stylized_target = glow(content_images,style_img)
            stylized_target_out = F.interpolate(stylized_target, (patch_H, patch_W), mode='bilinear', align_corners=False)
            stylized_target_out = torch.squeeze(stylized_target_out,0).permute(1,2,0)
            stylized_target_out = stylized_target_out.reshape(-1,3)


            alphainv_last_temp = alphainv_last.reshape(-1,3)
            one_temp = torch.ones_like(stylized_target_out)
            stylized_target_out = torch.where(alphainv_last_temp == 1.0, one_temp, stylized_target_out)


            target = stylized_target_out.flatten(0,-2)
            rgb_loss = cfg_train.weight_main * utils.img2mse(render_result['rgb_marched'], target)
            psnr = utils.mse2psnr(rgb_loss.detach())
            loss = rgb_loss
            
            
            if cfg_train.weight_entropy_last > 0:
                pout = render_result['alphainv_last'].clamp(1e-6, 1-1e-6)
                entropy_last_loss = -(pout*torch.log(pout) + (1-pout)*torch.log(1-pout)).mean()
                loss += cfg_train.weight_entropy_last * entropy_last_loss
            if cfg_train.weight_rgbper > 0:
                rgbper = (render_result['raw_rgb'] - target[render_result['ray_id']]).pow(2).sum(-1)
                rgbper_loss = (rgbper * render_result['weights'].detach()).sum() / len(rays_o)
                loss += cfg_train.weight_rgbper * rgbper_loss


        losses.update(loss.item())

        loss.backward()


        optimizer.step()
        psnr_lst.append(psnr.item())

        decay_steps = cfg_train.lrate_decay * 1000
        decay_factor = 0.1 ** (1/decay_steps)
        for i_opt_g, param_group in enumerate(optimizer.param_groups):
            param_group['lr'] = param_group['lr'] * decay_factor

        lr_temp = optimizer.state_dict()['param_groups'][0]['lr']
        # check log & save
        if global_step%args.i_print==0:
            eps_time = time.time() - time0
            eps_time_str = f'{eps_time//3600:02.0f}:{eps_time//60%60:02.0f}:{eps_time%60:02.0f}'

            tqdm.write(f'scene_rep_reconstruction ({stage}): iter {global_step:6d} / '
                       f'Loss: {losses.avg:.9f} / PSNR: {np.mean(psnr_lst):5.2f} / '
                       f'c_losses: {c_losses.avg:.6f} / '
                       f's_losses: {s_losses.avg:.6f} / '
                       f'lr: {lr_temp:.6f} / '
                       f'Eps: {eps_time_str}')
            psnr_lst = []

        if global_step%args.i_weights==0:
            path = os.path.join(cfg.basedir, cfg.expname, f'{stage}_{global_step:06d}.tar')
            torch.save({
                'global_step': global_step,
                'model_kwargs': model.get_kwargs(),
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
            }, path)
            print(f'scene_rep_reconstruction ({stage}): saved checkpoints at', path)

    if global_step != -1:
        torch.save({
            'global_step': global_step,
            'model_kwargs': model.get_kwargs(),
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, last_ckpt_path)
        print(f'scene_rep_reconstruction ({stage}): saved checkpoints at', last_ckpt_path)

def scene_rep_reconstruction(args, cfg, cfg_model, cfg_train, xyz_min, xyz_max, data_dict, stage, coarse_ckpt_path=None):

    # init
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if abs(cfg_model.world_bound_scale - 1) > 1e-9:
        xyz_shift = (xyz_max - xyz_min) * (cfg_model.world_bound_scale - 1) / 2
        xyz_min -= xyz_shift
        xyz_max += xyz_shift
    HW, Ks, near, far, i_train, i_val, i_test, poses, render_poses, images = [
        data_dict[k] for k in [
            'HW', 'Ks', 'near', 'far', 'i_train', 'i_val', 'i_test', 'poses', 'render_poses', 'images'
        ]
    ]

    # find whether there is existing checkpoint path
    last_ckpt_path = os.path.join(cfg.basedir, cfg.expname, f'{stage}_last.tar')
    if args.no_reload:
        reload_ckpt_path = None
    elif args.ft_path:
        reload_ckpt_path = args.ft_path
    elif os.path.isfile(last_ckpt_path):
        reload_ckpt_path = last_ckpt_path
    else:
        reload_ckpt_path = None

    # init model and optimizer
    if reload_ckpt_path is None:
        print(f'scene_rep_reconstruction ({stage}): train from scratch')
        start = 0
        # init model
        model_kwargs = copy.deepcopy(cfg_model)
        if cfg.data.ndc:
            print(f'scene_rep_reconstruction ({stage}): \033[96muse multiplane images\033[0m')
            num_voxels = model_kwargs.pop('num_voxels')
            if len(cfg_train.pg_scale) and reload_ckpt_path is None:
                num_voxels = int(num_voxels / (2**len(cfg_train.pg_scale)))
            model = upstmpi.UPSTMPI_DirectGO(
                xyz_min=xyz_min, xyz_max=xyz_max,
                num_voxels=num_voxels,
                mask_cache_path=coarse_ckpt_path,
                **model_kwargs)
        else:
            print(f'scene_rep_reconstruction ({stage}): \033[96muse dense voxel grid\033[0m')
            num_voxels = model_kwargs.pop('num_voxels')
            if len(cfg_train.pg_scale) and reload_ckpt_path is None:
                num_voxels = int(num_voxels / (2**len(cfg_train.pg_scale)))
            model = upst.UPST_DirectVoxGO(
                xyz_min=xyz_min, xyz_max=xyz_max,
                num_voxels=num_voxels,
                mask_cache_path=coarse_ckpt_path,
                **model_kwargs)

            if cfg_model.maskout_near_cam_vox:
                model.maskout_near_cam_vox(poses[i_train,:3,3], near)
        model = model.to(device)
        optimizer = utils.create_optimizer_or_freeze_model(model, cfg_train, global_step=0)
    else:
        print(f'scene_rep_reconstruction ({stage}): reload from {reload_ckpt_path}')
        if cfg.data.ndc:
            model_class = upstmpi.UPSTMPI_DirectGO
        else:
            model_class = upst.UPST_DirectVoxGO
        model = utils.load_model(model_class, reload_ckpt_path).to(device)
        optimizer = utils.create_optimizer_or_freeze_model(model, cfg_train, global_step=0)
        model, optimizer, start = utils.load_checkpoint(
                model, optimizer, reload_ckpt_path, args.no_reload_optimizer)

    # init rendering setup
    render_kwargs = {
        'near': data_dict['near'],
        'far': data_dict['far'],
        'bg': 1 if cfg.data.white_bkgd else 0,
        'stepsize': cfg_model.stepsize,
        'inverse_y': cfg.data.inverse_y,
        'flip_x': cfg.data.flip_x,
        'flip_y': cfg.data.flip_y,
    }

    # init batch rays sampler
    def gather_training_rays():
        if data_dict['irregular_shape']:
            rgb_tr_ori = [images[i].to('cpu' if cfg.data.load2gpu_on_the_fly else device) for i in i_train]
            if cfg_train.ray_sampler == 'in_maskcache':
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays_in_maskcache_sampling(
                        rgb_tr_ori=rgb_tr_ori,
                        train_poses=poses[i_train],
                        HW=HW[i_train], Ks=Ks[i_train],
                        ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                        flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y,
                        model=model, render_kwargs=render_kwargs)
            elif cfg_train.ray_sampler == 'flatten':
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays_flatten(
                    rgb_tr_ori=rgb_tr_ori,
                    train_poses=poses[i_train],
                    HW=HW[i_train], Ks=Ks[i_train], ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                    flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
            else:
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays_irregular_shape(
                    rgb_tr=rgb_tr_ori,
                    train_poses=poses[i_train],
                    HW=HW[i_train], Ks=Ks[i_train], ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                    flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)


        else:
            rgb_tr_ori = images[i_train].to('cpu' if cfg.data.load2gpu_on_the_fly else device)

            if cfg_train.ray_sampler == 'in_maskcache':
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays_in_maskcache_sampling(
                        rgb_tr_ori=rgb_tr_ori,
                        train_poses=poses[i_train],
                        HW=HW[i_train], Ks=Ks[i_train],
                        ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                        flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y,
                        model=model, render_kwargs=render_kwargs)
            elif cfg_train.ray_sampler == 'flatten':
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays_flatten(
                    rgb_tr_ori=rgb_tr_ori,
                    train_poses=poses[i_train],
                    HW=HW[i_train], Ks=Ks[i_train], ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                    flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)
            else:
                rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz = upst.get_training_rays(
                    rgb_tr=rgb_tr_ori,
                    train_poses=poses[i_train],
                    HW=HW[i_train], Ks=Ks[i_train], ndc=cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                    flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y)

        index_generator = upst.batch_indices_generator(len(rgb_tr), cfg_train.N_rand)
        batch_index_sampler = lambda: next(index_generator)
        return rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz, batch_index_sampler

    rgb_tr, rays_o_tr, rays_d_tr, viewdirs_tr, imsz, batch_index_sampler = gather_training_rays()

    # view-count-based learning rate
    if cfg_train.pervoxel_lr:
        def per_voxel_init():
            cnt = model.voxel_count_views(
                    rays_o_tr=rays_o_tr, rays_d_tr=rays_d_tr, imsz=imsz, near=near, far=far,
                    stepsize=cfg_model.stepsize, downrate=cfg_train.pervoxel_lr_downrate,
                    irregular_shape=data_dict['irregular_shape'])
            optimizer.set_pervoxel_lr(cnt)
            with torch.no_grad():
                model.density[cnt <= 2] = -100
        per_voxel_init()

    # GOGO
    torch.cuda.empty_cache()
    psnr_lst = []
    time0 = time.time()
    global_step = -1
    for global_step in trange(1+start, 1+cfg_train.N_iters):

        # renew occupancy grid
        if model.mask_cache is not None and (global_step + 500) % 1000 == 0:
            self_alpha = F.max_pool3d(model.activate_density(model.density), kernel_size=3, padding=1, stride=1)[0,0]
            model.mask_cache.mask &= (self_alpha > model.fast_color_thres)

        # progress scaling checkpoint
        if global_step in cfg_train.pg_scale:
            n_rest_scales = len(cfg_train.pg_scale)-cfg_train.pg_scale.index(global_step)-1
            cur_voxels = int(cfg_model.num_voxels / (2**n_rest_scales))
            if isinstance(model, upst.UPST_DirectVoxGO):
                model.scale_volume_grid(cur_voxels)
            elif isinstance(model, upstmpi.UPSTMPI_DirectGO):
                model.scale_volume_grid(cur_voxels, model.mpi_depth)
            else:
                raise NotImplementedErrore 
            optimizer = utils.create_optimizer_or_freeze_model(model, cfg_train, global_step=0)
            model.density.data.sub_(1)

        # random sample rays
        if cfg_train.ray_sampler in ['flatten', 'in_maskcache']:
            sel_i = batch_index_sampler()
            target = rgb_tr[sel_i]
            rays_o = rays_o_tr[sel_i]
            rays_d = rays_d_tr[sel_i]
            viewdirs = viewdirs_tr[sel_i]
        elif cfg_train.ray_sampler == 'random':
            sel_b = torch.randint(rgb_tr.shape[0], [cfg_train.N_rand])
            sel_r = torch.randint(rgb_tr.shape[1], [cfg_train.N_rand])
            sel_c = torch.randint(rgb_tr.shape[2], [cfg_train.N_rand])
            target = rgb_tr[sel_b, sel_r, sel_c]
            rays_o = rays_o_tr[sel_b, sel_r, sel_c]
            rays_d = rays_d_tr[sel_b, sel_r, sel_c]
            viewdirs = viewdirs_tr[sel_b, sel_r, sel_c]
        else:
            raise NotImplementedError

        if cfg.data.load2gpu_on_the_fly:
            target = target.to(device)
            rays_o = rays_o.to(device)
            rays_d = rays_d.to(device)
            viewdirs = viewdirs.to(device)

        # volume rendering
        # style_embedded = style_images
        style_embedded = None
        # render_result = model(rays_o, rays_d, viewdirs,torch.unsqueeze(style_embedded, 0), global_step=global_step, **render_kwargs)
        render_result = model(rays_o, rays_d, viewdirs,style_embedded, global_step=global_step, **render_kwargs)
        # gradient descent step
        optimizer.zero_grad(set_to_none=True)
        loss = cfg_train.weight_main * F.mse_loss(render_result['rgb_marched'], target)
        psnr = utils.mse2psnr(loss.detach())
        if cfg_train.weight_entropy_last > 0:
            pout = render_result['alphainv_last'].clamp(1e-6, 1-1e-6)
            entropy_last_loss = -(pout*torch.log(pout) + (1-pout)*torch.log(1-pout)).mean()
            loss += cfg_train.weight_entropy_last * entropy_last_loss
        if cfg_train.weight_rgbper > 0:
            rgbper = (render_result['raw_rgb'] - target[render_result['ray_id']]).pow(2).sum(-1)
            rgbper_loss = (rgbper * render_result['weights'].detach()).sum() / len(rays_o)
            loss += cfg_train.weight_rgbper * rgbper_loss
        loss.backward()

        if global_step<cfg_train.tv_before and global_step>cfg_train.tv_after and global_step%cfg_train.tv_every==0:
            if cfg_train.weight_tv_density>0:
                model.density_total_variation_add_grad(
                    cfg_train.weight_tv_density/len(rays_o), global_step<cfg_train.tv_dense_before)
            if cfg_train.weight_tv_k0>0:
                model.k0_total_variation_add_grad(
                    cfg_train.weight_tv_k0/len(rays_o), global_step<cfg_train.tv_dense_before)

        optimizer.step()
        psnr_lst.append(psnr.item())

        # update lr
        decay_steps = cfg_train.lrate_decay * 1000
        decay_factor = 0.1 ** (1/decay_steps)
        for i_opt_g, param_group in enumerate(optimizer.param_groups):
            param_group['lr'] = param_group['lr'] * decay_factor

        # check log & save
        if global_step%args.i_print==0:
            eps_time = time.time() - time0
            eps_time_str = f'{eps_time//3600:02.0f}:{eps_time//60%60:02.0f}:{eps_time%60:02.0f}'
            tqdm.write(f'scene_rep_reconstruction ({stage}): iter {global_step:6d} / '
                       f'Loss: {loss.item():.9f} / PSNR: {np.mean(psnr_lst):5.2f} / '
                       f'Eps: {eps_time_str}')
            psnr_lst = []

        if global_step%args.i_weights==0:
            path = os.path.join(cfg.basedir, cfg.expname, f'{stage}_{global_step:06d}.tar')
            torch.save({
                'global_step': global_step,
                'model_kwargs': model.get_kwargs(),
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
            }, path)
            print(f'scene_rep_reconstruction ({stage}): saved checkpoints at', path)

    if global_step != -1:
        torch.save({
            'global_step': global_step,
            'model_kwargs': model.get_kwargs(),
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, last_ckpt_path)
        print(f'scene_rep_reconstruction ({stage}): saved checkpoints at', last_ckpt_path)



def train(args, cfg, data_dict):

    # init
    print('train: start')
    eps_time = time.time()
    os.makedirs(os.path.join(cfg.basedir, cfg.expname), exist_ok=True)
    with open(os.path.join(cfg.basedir, cfg.expname, 'args.txt'), 'w') as file:
        for arg in sorted(vars(args)):
            attr = getattr(args, arg)
            file.write('{} = {}\n'.format(arg, attr))
    cfg.dump(os.path.join(cfg.basedir, cfg.expname, 'config.py'))

    # coarse geometry searching
    eps_coarse = time.time()
    xyz_min_coarse, xyz_max_coarse = compute_bbox_by_cam_frustrm(args=args, cfg=cfg, **data_dict)
    if cfg.coarse_train.N_iters > 0:
        scene_rep_reconstruction(
                args=args, cfg=cfg,
                cfg_model=cfg.coarse_model_and_render, cfg_train=cfg.coarse_train,
                xyz_min=xyz_min_coarse, xyz_max=xyz_max_coarse,
                data_dict=data_dict, stage='coarse')
        eps_coarse = time.time() - eps_coarse
        eps_time_str = f'{eps_coarse//3600:02.0f}:{eps_coarse//60%60:02.0f}:{eps_coarse%60:02.0f}'
        print('train: coarse geometry searching in', eps_time_str)
        coarse_ckpt_path = os.path.join(cfg.basedir, cfg.expname, f'coarse_last.tar')
    else:
        print('train: skip coarse geometry searching')
        coarse_ckpt_path = None

    # fine detail reconstruction
    eps_fine = time.time()
    if cfg.data.ndc:
        xyz_min_fine, xyz_max_fine = xyz_min_coarse.clone(), xyz_max_coarse.clone()
    else:
        xyz_min_fine, xyz_max_fine = compute_bbox_by_coarse_geo(
                model_class=upst.UPST_DirectVoxGO, model_path=coarse_ckpt_path,
                thres=cfg.fine_model_and_render.bbox_thres)
    scene_rep_reconstruction(
            args=args, cfg=cfg,
            cfg_model=cfg.fine_model_and_render, cfg_train=cfg.fine_train,
            xyz_min=xyz_min_fine, xyz_max=xyz_max_fine,
            data_dict=data_dict, stage='fine',
            coarse_ckpt_path=coarse_ckpt_path)
    eps_fine = time.time() - eps_fine
    eps_time_str = f'{eps_fine//3600:02.0f}:{eps_fine//60%60:02.0f}:{eps_fine%60:02.0f}'
    print('train: fine detail reconstruction in', eps_time_str)



    # Style transfer and reconstruction
    coarse_ckpt_path = os.path.join(cfg.basedir, cfg.expname, f'fine_last.tar')
    eps_fine_style = time.time()
    scene_rep_reconstruction_style(
            args=args, cfg=cfg,
            cfg_model=cfg.fine_model_and_render, cfg_train=cfg.style_train,
            data_dict=data_dict, stage='fine_style',
            coarse_ckpt_path=coarse_ckpt_path)

    eps_fine_style = time.time() - eps_fine_style
    eps_time_str = f'{eps_fine_style//3600:02.0f}:{eps_fine_style//60%60:02.0f}:{eps_fine_style%60:02.0f}'
    print('train: style reconstruction in', eps_time_str)


    eps_time = time.time() - eps_time
    eps_time_str = f'{eps_time//3600:02.0f}:{eps_time//60%60:02.0f}:{eps_time%60:02.0f}'
    print('train: finish (eps time', eps_time_str, ')')


if __name__=='__main__':

    # load setup
    parser = config_parser()
    args = parser.parse_args()
    cfg = mmcv.Config.fromfile(args.config)

    # init enviroment
    if torch.cuda.is_available():
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')
    seed_everything()
    data_dict = load_everything(cfg=cfg)

    # export scene bbox and camera poses in 3d for debugging and visualization
    if args.export_bbox_and_cams_only:
        print('Export bbox and cameras...')
        xyz_min, xyz_max = compute_bbox_by_cam_frustrm(args=args, cfg=cfg, **data_dict)
        poses, HW, Ks, i_train = data_dict['poses'], data_dict['HW'], data_dict['Ks'], data_dict['i_train']
        near, far = data_dict['near'], data_dict['far']
        cam_lst = []
        for c2w, (H, W), K in zip(poses[i_train], HW[i_train], Ks[i_train]):
            rays_o, rays_d, viewdirs = upst.get_rays_of_a_view(
                    H, W, K, c2w, cfg.data.ndc, inverse_y=cfg.data.inverse_y,
                    flip_x=cfg.data.flip_x, flip_y=cfg.data.flip_y,)
            cam_o = rays_o[0,0].cpu().numpy()
            cam_d = rays_d[[0,0,-1,-1],[0,-1,0,-1]].cpu().numpy()
            cam_lst.append(np.array([cam_o, *(cam_o+cam_d*max(near, far*0.05))]))
        np.savez_compressed(args.export_bbox_and_cams_only,
            xyz_min=xyz_min.cpu().numpy(), xyz_max=xyz_max.cpu().numpy(),
            cam_lst=np.array(cam_lst))
        print('done')
        sys.exit()

    if args.export_coarse_only:
        print('Export coarse visualization...')
        with torch.no_grad():
            ckpt_path = os.path.join(cfg.basedir, cfg.expname, 'coarse_last.tar')
            model = utils.load_model(upst.UPST_DirectVoxGO, ckpt_path).to(device)
            alpha = model.activate_density(model.density).squeeze().cpu().numpy()
            rgb = torch.sigmoid(model.k0).squeeze().permute(1,2,3,0).cpu().numpy()
        np.savez_compressed(args.export_coarse_only, alpha=alpha, rgb=rgb)
        print('done')
        sys.exit()

    # train
    if not args.render_only:
        train(args, cfg, data_dict)
    # load model for rendring



    if args.render_geometry:
        if args.render_test or args.render_train or args.render_video:
            if args.ft_path:
                ckpt_path = args.ft_path
            else:
                ckpt_path = os.path.join(cfg.basedir, cfg.expname, 'fine_last.tar')
            ckpt_name = ckpt_path.split('/')[-1][:-4]
            if cfg.data.ndc:
                model_class = upstmpi.UPSTMPI_DirectGO
            else:
                model_class = upst.UPST_DirectVoxGO
            model = utils.load_model(model_class, ckpt_path).to(device)
            stepsize = cfg.fine_model_and_render.stepsize
            render_viewpoints_kwargs = {
                'model': model,
                'ndc': cfg.data.ndc,
                'render_kwargs': {
                    'near': data_dict['near'],
                    'far': data_dict['far'],
                    'bg': 1 if cfg.data.white_bkgd else 0,
                    'stepsize': stepsize,
                    'inverse_y': cfg.data.inverse_y,
                    'flip_x': cfg.data.flip_x,
                    'flip_y': cfg.data.flip_y,
                    'render_depth': True,
                },
            }

        # render trainset and eval
        if args.render_train:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_train_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints(
                    render_poses=data_dict['poses'][data_dict['i_train']],
                    HW=data_dict['HW'][data_dict['i_train']],
                    Ks=data_dict['Ks'][data_dict['i_train']],
                    gt_imgs=[data_dict['images'][i].cpu().numpy() for i in data_dict['i_train']],
                    savedir=testsavedir,
                    eval_ssim=args.eval_ssim, eval_lpips_alex=args.eval_lpips_alex, eval_lpips_vgg=args.eval_lpips_vgg,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)

        # render testset and eval
        if args.render_test:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_test_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints(
                    render_poses=data_dict['poses'][data_dict['i_test']],
                    HW=data_dict['HW'][data_dict['i_test']],
                    Ks=data_dict['Ks'][data_dict['i_test']],
                    gt_imgs=[data_dict['images'][i].cpu().numpy() for i in data_dict['i_test']],
                    savedir=testsavedir,
                    eval_ssim=args.eval_ssim, eval_lpips_alex=args.eval_lpips_alex, eval_lpips_vgg=args.eval_lpips_vgg,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)

        # render video
        if args.render_video:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_video_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints(
                    render_poses=data_dict['render_poses'],
                    HW=data_dict['HW'][data_dict['i_test']][[0]].repeat(len(data_dict['render_poses']), 0),
                    Ks=data_dict['Ks'][data_dict['i_test']][[0]].repeat(len(data_dict['render_poses']), 0),
                    render_factor=args.render_video_factor,
                    savedir=testsavedir,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)


    if args.render_style:
        if args.render_test or args.render_train or args.render_video:
            if args.ft_path:
                ckpt_path = args.ft_path
            else:
                ckpt_path = os.path.join(cfg.basedir, cfg.expname, 'fine_style_last.tar')
            ckpt_name = ckpt_path.split('/')[-1][:-4]
            if cfg.data.ndc:
                model_class = upstmpi.UPSTMPI_DirectGO
            else:
                model_class = upst.UPST_DirectVoxGO
            model = utils.load_model(model_class, ckpt_path).to(device)
            stepsize = cfg.fine_model_and_render.stepsize
            render_viewpoints_kwargs = {
                'model': model,
                'ndc': cfg.data.ndc,
                'render_kwargs': {
                    'near': data_dict['near'],
                    'far': data_dict['far'],
                    'bg': 1 if cfg.data.white_bkgd else 0,
                    'stepsize': stepsize,
                    'inverse_y': cfg.data.inverse_y,
                    'flip_x': cfg.data.flip_x,
                    'flip_y': cfg.data.flip_y,
                    'render_depth': True,
                },
            }

        style_img = args.style_img
        # render trainset and eval
        if args.render_train:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_train_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints_style(
                    style_img=style_img,
                    render_poses=data_dict['poses'][data_dict['i_train']],
                    HW=data_dict['HW'][data_dict['i_train']],
                    Ks=data_dict['Ks'][data_dict['i_train']],
                    gt_imgs=[data_dict['images'][i].cpu().numpy() for i in data_dict['i_train']],
                    savedir=testsavedir,
                    eval_ssim=args.eval_ssim, eval_lpips_alex=args.eval_lpips_alex, eval_lpips_vgg=args.eval_lpips_vgg,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)

        # render testset and eval
        if args.render_test:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_test_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints_style(
                    style_img=style_img,
                    render_poses=data_dict['poses'][data_dict['i_test']],
                    HW=data_dict['HW'][data_dict['i_test']],
                    Ks=data_dict['Ks'][data_dict['i_test']],
                    gt_imgs=[data_dict['images'][i].cpu().numpy() for i in data_dict['i_test']],
                    savedir=testsavedir,
                    eval_ssim=args.eval_ssim, eval_lpips_alex=args.eval_lpips_alex, eval_lpips_vgg=args.eval_lpips_vgg,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)

        # render video
        if args.render_video:
            testsavedir = os.path.join(cfg.basedir, cfg.expname, f'render_video_{ckpt_name}')
            os.makedirs(testsavedir, exist_ok=True)
            rgbs, depths = render_viewpoints_style(
                    style_img=style_img,
                    render_poses=data_dict['render_poses'],
                    HW=data_dict['HW'][data_dict['i_test']][[0]].repeat(len(data_dict['render_poses']), 0),
                    Ks=data_dict['Ks'][data_dict['i_test']][[0]].repeat(len(data_dict['render_poses']), 0),
                    render_factor=args.render_video_factor,
                    savedir=testsavedir,
                    **render_viewpoints_kwargs)
            imageio.mimwrite(os.path.join(testsavedir, 'video.rgb.mp4'), utils.to8b(rgbs), fps=30, quality=8)
            imageio.mimwrite(os.path.join(testsavedir, 'video.depth.mp4'), utils.to8b(1 - depths / np.max(depths)), fps=30, quality=8)



    print('Done')