train_style.py

# Copyright (c) Meta Platforms, Inc. and affiliates.

import datetime
import os
import sys
from typing import List
import io
import imageio
import numpy as np
import torch
torch.backends.cudnn.enabled = False
import torch.nn as nn
import torch.nn.functional as F
from torchvision.utils import make_grid
from easydict import EasyDict as edict
from torch.utils.tensorboard import SummaryWriter
from tqdm.auto import tqdm
import matplotlib
from PIL import Image, ImageFile
from tqdm import trange

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from torch_efficient_distloss import (
    eff_distloss,
    eff_distloss_native,
    flatten_eff_distloss,
)
from utils import visualize_depth_numpy


from camera import (
    pose_to_mtx,
    cam2world,
    lie,
    pose,
    procrustes_analysis,
    rotation_distance,
    get_novel_view_poses,
)

from dataLoader.wikiart import getDataLoader as getStyleDataLoader
from dataLoader import dataset_dict
from dataLoader.ray_utils import (
    get_ray_directions_blender,
    get_ray_directions_lean,
    get_rays,
    get_rays_lean,
    get_rays_with_batch,
    ndc_rays_blender,
    ndc_rays_blender2,
    normalize_vgg, denormalize_vgg
)
# from models.tensoRF import TensorVMSplit, TensorVMSplit_TimeEmbedding
# from models.tensorStyRF import TensorVMSplit, TensorVMSplit_TimeEmbedding
# from models.tensorWCTRF import TensorVMSplit, TensorVMSplit_TimeEmbedding
from models.tensorStyDyRF import TensorVMSplit, TensorVMSplit_TimeEmbedding
from models.vggNetworks import encoder3, decoder3
from models.vggNetworks import encoder4, decoder4
from models.vggNetworks import encoder5
# from models.wct_matrix import MulLayer
# from models.wct_matrix3d import MulLayer as MulLayer3d
from models.wct_matrix4d import MulLayer as MulLayer4d
import models.wct_cspn as model_spn
import models.wct_update_model as update_model
from models.wct_criterion import LossCriterion
from models.wct_criterion import CorrelationLoss
from models.wct_criterion import GradLoss
from models.wct_criterion import TVLoss as WCT_TV_loss
from models.temporal_criterion import TemporalLoss
from models.linearWCT.build_model import LinearWCT
from opt import config_parser
from renderer import (
    evaluation,
    evaluation_path,
    OctreeRender_trilinear_fast,
    render,
    induce_flow,
    render_3d_point,
    render_single_3d_point,
    NDC2world,
    induce_flow_single,
    raw2outputs,
    sampleXYZ,
    contract2world,
    raw2outputs_feature,
    raw2outputs_feature_only
)
from utils import cal_n_samples, convert_sdf_samples_to_ply, N_to_reso, TVLoss
from flow_viz import flow_to_image

# # 限制cpu core使用数
# import os
# from multiprocessing import cpu_count

# # cpu_num = cpu_count() // 2 # 自动获取最大核心数目
# cpu_num = int(cpu_count() * opt.cpu_percentage)
# os.environ ['OMP_NUM_THREADS'] = str(cpu_num)
# os.environ ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
# os.environ ['MKL_NUM_THREADS'] = str(cpu_num)
# os.environ ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
# os.environ ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
# torch.set_num_threads(cpu_num)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

renderer = OctreeRender_trilinear_fast


# Dummy tensorboard logger
class DummyWriter:
    def add_scalar(*args, **kwargs):
        pass

    def add_images(*args, **kwargs):
        pass


class SimpleSampler:
    def __init__(self, total, batch):
        self.total = total
        self.batch = batch
        self.curr = total
        self.ids = None

    def nextids(self):
        self.curr += self.batch
        if self.curr + self.batch > self.total:
            self.ids = torch.LongTensor(np.random.permutation(self.total))
            self.curr = 0
        return self.ids[self.curr : self.curr + self.batch]


def InfiniteSampler(n):
    # i = 0
    i = n - 1
    order = np.random.permutation(n)
    while True:
        yield order[i]
        i += 1
        if i >= n:
            np.random.seed()
            order = np.random.permutation(n)
            i = 0

class InfiniteSamplerWrapper(torch.utils.data.sampler.Sampler):
    def __init__(self, num_samples):
        self.num_samples = num_samples

    def __iter__(self):
        return iter(InfiniteSampler(self.num_samples))

    def __len__(self):
        return 2 ** 31


def ids2pixel(W, H, ids):
    """
    Regress pixel coordinates from
    """
    col = ids % W
    row = (ids // W) % H
    view_ids = ids // (W * H)
    return col, row, view_ids


def reconstruction(args):
    # init style dataset
    # Note: this code is borrowed from stylerf
    Image.MAX_IMAGE_PIXELS = None  # Disable DecompressionBombError
    ImageFile.LOAD_TRUNCATED_IMAGES = True # Disable OSError: image file is truncated
    style_loader = getStyleDataLoader(args.wikiartdir, batch_size=4, sampler=InfiniteSamplerWrapper, 
                    image_side_length=256, num_workers=4)
    style_iter = iter(style_loader)

    # init dataset
    dataset = dataset_dict[args.dataset_name]
    train_dataset = dataset(
        args.datadir,
        split="train",
        downsample=args.downsample_train,
        is_stack=True,
        use_disp=args.use_disp,
        use_foreground_mask=args.use_foreground_mask,
        with_GT_poses=args.with_GT_poses,
        ray_type=args.ray_type,
    )
    white_bg = train_dataset.white_bg
    near_far = train_dataset.near_far
    W, H = train_dataset.img_wh

    # init resolution
    upsamp_list = args.upsamp_list
    n_lamb_sigma = args.n_lamb_sigma
    n_lamb_sh = args.n_lamb_sh

    if args.add_timestamp:
        logfolder = f'{args.basedir}/{args.expname}{datetime.datetime.now().strftime("-%Y%m%d-%H%M%S")}'
    else:
        logfolder = f"{args.basedir}/{args.expname}"

    # init log fileinit log file
    os.makedirs(logfolder, exist_ok=True)
    os.makedirs(f"{logfolder}/imgs_vis", exist_ok=True)
    os.makedirs(f"{logfolder}/imgs_rgba", exist_ok=True)
    os.makedirs(f"{logfolder}/rgba", exist_ok=True)
    summary_writer = SummaryWriter(logfolder)

    # init parameters
    aabb = train_dataset.scene_bbox.to(device)
    reso_cur = N_to_reso(args.N_voxel_init, aabb)
    nSamples = min(args.nSamples, cal_n_samples(reso_cur, args.step_ratio))

    # 导入模型参数
    # dynamic
    ckpt = torch.load(args.ckpt_feature, map_location=device)
    kwargs = ckpt["kwargs"]
    poses_mtx = kwargs.pop("se3_poses").to(device)
    focal_refine = kwargs.pop("focal_ratio_refine").to(device)
    kwargs.update({"device": device})
    tensorf = eval(args.model_name)(**kwargs)
    # 调整模型设置为feature模式
    tensorf.change_to_feature_mod(args.n_lamb_sh, device)
    tensorf.load(ckpt)
    # static
    ckpt_static = torch.load(args.ckpt_feature[:-3] + "_static.th", map_location=device)
    kwargs_static = ckpt_static["kwargs"]
    poses_mtx = kwargs_static.pop("se3_poses").to(device)
    focal_refine = kwargs_static.pop("focal_ratio_refine").to(device)
    kwargs_static.update({"device": device})
    tensorf_static = TensorVMSplit(**kwargs_static)
    # 调整模型设置为feature模式
    tensorf_static.change_to_feature_mod(args.n_lamb_sh, device)
    tensorf_static.load(ckpt_static)

    # 调整模型设置为style模式
    tensorf.change_to_style_mod(device)
    tensorf_static.change_to_style_mod(device)

    vgg = encoder3()
    dec = decoder3()
    vgg5 = encoder5()
    vgg_dir = "pretrained/vgg_r31.pth"
    decoder_dir = "pretrained/dec_r31.pth"
    loss_network_dir = "pretrained/vgg_r51.pth"
    vgg.load_state_dict(torch.load(vgg_dir))
    dec.load_state_dict(torch.load(decoder_dir))
    vgg5.load_state_dict(torch.load(loss_network_dir))
    # 不参与训练
    for param in vgg.parameters():
        param.requires_grad = False
    for param in dec.parameters():
        param.requires_grad = False
    for param in vgg5.parameters():
        param.requires_grad = False
    vgg = vgg.to(device)
    dec = dec.to(device)
    vgg5 = vgg5.to(device)

    # 导入linear wct模型
    print("==> Loading linear wct")
    linear_wct = LinearWCT(
                layer='r41', 
                vgg_dir="pretrained/vgg_r41.pth", 
                decoder_dir="pretrained/dec_r41.pth",
                matrixPath="pretrained/r41.pth", 
                spn_dir="pretrained/r41_spn.pth", 
                device=device)

    # matrix = MulLayer("r31").to(device)
    # matrix = MulLayer3d("r31").to(device)
    matrix = MulLayer4d("r31").to(device)
    spn = model_spn.resnet50(pretrained = False).to(device)
    
    style_layers = ["r11", "r21", "r31", "r41"]
    content_layers = ["r41"]
    style_weight = 0.02
    content_weight = 1.0
    criterion = LossCriterion(style_layers,
                          content_layers,
                          style_weight,
                          content_weight)

    criterion_corr = CorrelationLoss()
    criterion_grad = GradLoss()
    criterion_contentProp = nn.MSELoss()
    criterion_tv = WCT_TV_loss()
    criterion_temporal = TemporalLoss()

    grad_vars = list(matrix.parameters()) + list(spn.parameters())
    
    if args.lr_decay_iters > 0:
        lr_factor = args.lr_decay_target_ratio ** (1 / args.lr_decay_iters)
    else:
        args.lr_decay_iters = args.n_iters
        lr_factor = args.lr_decay_target_ratio ** (1 / args.n_iters)

    print("lr decay", args.lr_decay_target_ratio, args.lr_decay_iters)

    optimizer = torch.optim.Adam(grad_vars, betas=(0.9, 0.99))

    # 数据集所有ray的数据
    allrgbs = train_dataset.all_rgbs
    # allfeatures = train_dataset.all_features
    allts = train_dataset.all_ts
    if args.with_GT_poses:
        allposes = train_dataset.all_poses  # (12, 3, 4)

    # ray sampling
    W, H = train_dataset.img_wh
    directions = get_ray_directions_blender(
        H, W, [focal_refine.cpu(), focal_refine.cpu()]
    ).to(poses_mtx.device)  # (H, W, 3)
    all_rays = []
    for i in range(poses_mtx.shape[0]):
        c2w = poses_mtx[i]
        rays_o, rays_d = get_rays(directions, c2w)  # both (h*w, 3)
        if args.ray_type == "ndc":
            rays_o, rays_d = ndc_rays_blender(
                H, W, focal_refine.cpu(), 1.0, rays_o, rays_d
            )
        all_rays += [torch.cat([rays_o, rays_d], 1)]  # (h*w, 6)
    all_rays = torch.stack(all_rays, 0).to(device)  # num_frames x h*w x 6
    if args.multiview_dataset:
        # duplicate poses for multiple time instances
        all_rays = torch.tile(all_rays, (args.N_voxel_t, 1, 1))

    print(f"allrgbs: {allrgbs.shape}")
    # print(f"allfeatures: {allfeatures.shape}")
    print(f"allts: {allts.shape}")
    print(f"all_rays: {all_rays.shape}")

    num_frames = all_rays.shape[0]
    all_rays_stack = all_rays.reshape(num_frames, H, W, -1)  # num_frames x height x width x 6
    all_rgbs_stack = allrgbs.reshape(num_frames, H, W, -1)  # num_frames x height x width x 3
    all_ts_stack = allts.reshape(num_frames, H, W)
    # all_features_stack = allfeatures.reshape(num_frames, H, W, -1)  # num_frames x height x width x 256
    
    all_rays = all_rays.reshape(num_frames*H*W, -1)
    all_rgbs = allrgbs.reshape(num_frames*H*W, -1)
    all_ts = allts.reshape(num_frames*H*W)
    # all_features = allfeatures.reshape(num_frames*H*W, -1)

    print('==> Extracting canonical feature volume')
    canonical_features_cache_path = os.path.join(logfolder, "canonical_features.pth")
    if not os.path.exists(canonical_features_cache_path):
        with torch.no_grad():
            # 采样-1到1空间上的点
            n_resolution = 32
            x = torch.linspace(-1, 1, n_resolution)  # 64
            xs, ys, zs = torch.meshgrid(x, x, x)
            # xs: [64, 64, 64]
            # ys: [64, 64, 64]
            # zs: [64, 64, 64]
            xyz_sampled = torch.cat([xs.unsqueeze(-1), ys.unsqueeze(-1), zs.unsqueeze(-1)], dim=0)  # [64, 64, 64, 3]
            xyz_sampled = xyz_sampled.view(-1, 3)  # [64 * 64 * 64, 3]

            # canonical_feature_volume = torch.zeros((xyz_sampled.shape[0], 256))

            chunk = args.batch_size  # 512
            N_rays_chunk = xyz_sampled.shape[0]
            canonical_features = []
            for chunk_idx in range(N_rays_chunk // chunk + int(N_rays_chunk % chunk > 0)):
                xyz_chunk = xyz_sampled[chunk_idx * chunk : (chunk_idx + 1) * chunk].to(device)  # [batch_size, 3]
                
                feature_chunk = tensorf.compute_canonical_feature(xyz_chunk)  # [batch_size, 256]

                canonical_features.append(feature_chunk)

            canonical_features = torch.cat(canonical_features, dim=0).detach()  # [64 * 64 * 64， 256]
            canonical_features = canonical_features.reshape(n_resolution, n_resolution, n_resolution, 256)  # [64 x 64 x 64 x 256]
            canonical_features = canonical_features.permute(3, 0, 1, 2).unsqueeze(dim=0)  # [1 x 256 x 64 x 64 x 64]
            
            canonical_features = canonical_features.to(device)  
            torch.save(canonical_features, canonical_features_cache_path)
    else:
        print(f'Loading from cache: {canonical_features_cache_path}')
        canonical_features = torch.load(canonical_features_cache_path)
    print(f"canonical_features: {canonical_features.shape}")

    # trainingSampler = SimpleSampler(allts.shape[0], args.batch_size)
    # trainingSampler_2 = SimpleSampler(allts.shape[0], args.batch_size)
    frameSampler = iter(InfiniteSamplerWrapper(train_dataset.num_images)) # every next(sampler) returns a frame index

    pbar = tqdm(
        range(args.n_iters), miniters=args.progress_refresh_rate, file=sys.stdout
    )
    for iteration in pbar:

        
        patch_size = args.patch_size
        # N_rays_all = rays.shape[0]
        N_samples = -1
        chunk = args.batch_size  # 512
        h_rays = H
        w_rays = W

        style_img = next(style_iter)[0].to(device)
        # breakpoint()

        frame_idx = next(frameSampler)
        start_h = np.random.randint(0, h_rays-patch_size+1)
        start_w = np.random.randint(0, w_rays-patch_size+1)
        if white_bg:
            # move random sampled patches into center
            mid_h, mid_w = (h_rays-patch_size+1)/2, (w_rays-patch_size+1)/2
            if mid_h-start_h>=1:
                start_h += np.random.randint(0, mid_h-start_h)
            elif mid_h-start_h<=-1:
                start_h += np.random.randint(mid_h-start_h, 0)
            if mid_w-start_w>=1:
                start_w += np.random.randint(0, mid_w-start_w)
            elif mid_w-start_w<=-1:
                start_w += np.random.randint(mid_w-start_w, 0)

        rays_patch = all_rays_stack[frame_idx, start_h:start_h+patch_size, 
                                                start_w:start_w+patch_size, :].reshape(-1, 6).to(device)
        # [patch*patch, 6]
        
        rgbs_patch = all_rgbs_stack[frame_idx, start_h:(start_h+patch_size), 
                                                start_w:(start_w+patch_size), :].to(device)

        ts_patch = all_ts_stack[frame_idx, start_h:start_h+patch_size, 
                                            start_w:start_w+patch_size].reshape(-1).to(device)

        # features_patch = all_features_stack[frame_idx, start_h:start_h+patch_size, 
        #                                                 start_w:start_w+patch_size, :].reshape(-1, 256).to(device)

        with torch.no_grad():                 
            N_rays_patch = rays_patch.shape[0]
            feature_map_list = []
            feature_map_d_list = []
            feature_map_s_list = []
            # feature_points_d_list = []
            # feature_points_s_list = []
            feature_points_list = []
            for chunk_idx in range(N_rays_patch // chunk + int(N_rays_patch % chunk > 0)):
                rays_chunk = rays_patch[chunk_idx * chunk : (chunk_idx + 1) * chunk].to(device)
                ts_chunk = ts_patch[chunk_idx * chunk : (chunk_idx + 1) * chunk].to(device)
                # features_chunk = features[chunk_idx * chunk : (chunk_idx + 1) * chunk].to(device)

                xyz_sampled, z_vals, ray_valid = sampleXYZ(
                    tensorf,
                    rays_chunk,
                    N_samples=N_samples,
                    ray_type=args.ray_type,
                    is_train=False,
                )
                # static 
                (
                    _, 
                    _, 
                    _, 
                    _, 
                    _, 
                    _, 
                    rgb_point_static, 
                    sigma_static, 
                    _, 
                    _,
                    feature_point_static, 
                ) = tensorf_static.render_rgb_feature_map(
                    rays_chunk,
                    ts_chunk,
                    None,
                    xyz_sampled,
                    z_vals,
                    ray_valid,
                    is_train=False,
                    white_bg=white_bg,
                    ray_type=args.ray_type,
                    N_samples=N_samples,
                )
                # dynamic
                (
                    _,
                    _,
                    blending,
                    pts_ref,
                    _,
                    _,
                    rgb_point_dynamic,
                    sigma_dynamic,
                    z_val_dynamic,
                    dist_dynamic,
                    feature_point_dynamic,
                ) = tensorf.render_rgb_feature_map(
                    rays_chunk,
                    ts_chunk,
                    None,
                    xyz_sampled,
                    z_vals,
                    ray_valid,
                    is_train=False,
                    white_bg=white_bg,
                    ray_type=args.ray_type,
                    N_samples=N_samples,
                )
                # blending
                (
                    feature_map_full,
                    depth_map_full,
                    acc_map_full,
                    weights_full,
                    feature_map_s,
                    depth_map_s,
                    acc_map_s,
                    weights_s,
                    feature_map_d,
                    depth_map_d,
                    acc_map_d,
                    weights_d,
                    dynamicness_map,
                ) = raw2outputs_feature_only(
                    feature_point_static,
                    sigma_static,
                    feature_point_dynamic,
                    sigma_dynamic,
                    dist_dynamic,
                    blending,
                    z_val_dynamic,
                    rays_chunk,
                    ray_type=args.ray_type,
                )

                feature_map_list.append(feature_map_full)
                feature_map_d_list.append(feature_map_d)
                feature_map_s_list.append(feature_map_s)
                # feature_points = blending[..., None] * feature_point_dynamic + (1.0 - blending[..., None]) * feature_point_static
                # feature_points = feature_points.permute(0, 2, 1)  # n_rays x 256 x n_samples
                # feature_points = F.interpolate(feature_points, size=(patch_size // 4), mode='linear')
                # feature_points = feature_points.permute(0, 2, 1)
                # feature_points_list.append(feature_points)  
                # feature_points_d_list.append(feature_point_dynamic)
                # feature_points_s_list.append(feature_point_static)

            feature_map_patch = torch.cat(feature_map_list)  # H*W x 256
            feature_map_d_patch = torch.cat(feature_map_d_list)  # H*W x 256
            feature_map_s_patch = torch.cat(feature_map_s_list)  # H*W x 256
            # feature_points_patch = torch.cat(feature_points_list)  # H*W x n_samples x 256

            feature_map_patch = feature_map_patch.reshape(patch_size, patch_size, 256)[None,...].permute(0,3,1,2)
            feature_map_d_patch = feature_map_d_patch.reshape(patch_size, patch_size, 256)[None,...].permute(0,3,1,2)
            feature_map_s_patch = feature_map_s_patch.reshape(patch_size, patch_size, 256)[None,...].permute(0,3,1,2)
            # 1 x 256 x 256 x 256
            # feature_n_samples = feature_points_patch.shape[1]
            # feature_points_patch = feature_points_patch.reshape(patch_size, patch_size, feature_n_samples, 256)[None,...].permute(0,4,1,2,3)

            feature_size = patch_size // 4
            feature_map_patch = F.interpolate(feature_map_patch, size=(feature_size, feature_size), mode='bilinear')
            feature_map_d_patch = F.interpolate(feature_map_d_patch, size=(feature_size, feature_size), mode='bilinear')
            feature_map_s_patch = F.interpolate(feature_map_s_patch, size=(feature_size, feature_size), mode='bilinear')
            # 1 x 256 x 64 x 64
            # volume_size = patch_size // 4
            # feature_points_patch = F.interpolate(feature_points_patch, size=(volume_size, volume_size, volume_size), mode='trilinear')
            # 1 x 256 x 64 x 64 x 64
            rgbs_train = rgbs_patch[None,...].permute(0,3,1,2)  # 1 x 3 x patch_size x patch_size
            content_rgb = rgbs_train

        torch.cuda.empty_cache()

        # 提取风格图像的特征
        sF = vgg(style_img)  # batch_size_sF x 256 x 64 x 64
        # 提取内容图像的特征，直接渲染特征图
        cF = feature_map_patch.detach()  # 1 x 256 x 64 x 64
        cF_d = feature_map_d_patch.detach()  # 1 x 256 x 64 x 64
        # cF_s = feature_map_s_patch.detach()  # 1 x 256 x 64 x 64
        # content_recon = dec(cF)

        cF = cF.repeat(sF.shape[0], 1, 1, 1)  # batch_size_sF x 256 x 64 x 64
        cF_d = cF_d.repeat(sF.shape[0], 1, 1, 1)  # batch_size_sF x 256 x 64 x 64
        content_rgb = content_rgb.repeat(sF.shape[0], 1, 1, 1)  # batch_size_sF x 3 x 256 x 256

        # WCT
        feature, transmatrix = matrix(canonical_features, cF, sF)
        feature_d, transmatrix_d = matrix(canonical_features, cF_d, sF)
        # feature_s, transmatrix_s = matrix(canonical_features, cF_s, sF)

        # 将风格迁移后的特征解码回原始rgb空间
        transfer = dec(feature)  # 1 x 3 x 256 x 256
        transfer_d = dec(feature_d)  # 1 x 3 x 256 x 256
        # transfer_s = dec(feature_s)  # 1 x 3 x 256 x 256

        # 去畸变网络，利用content image的信息来引导去噪过程
        propagated = spn(content_rgb, transfer)

        # recon_rgb = tensorf.decoder(feature_map_patch)
        # recon_rgb_d = tensorf.decoder(feature_map_d_patch)
        # recon_rgb_s = tensorf.decoder(feature_map_s_patch)
        
        # 重建回原图
        recon_rgb = dec(feature_map_patch)
        recon_rgb_d = dec(feature_map_d_patch)
        recon_rgb_s = dec(feature_map_s_patch)
        # # 1 x 3 x 256 x 256

        # 计算loss，提取vgg的多层级特征
        sF_loss = vgg5(style_img)
        cF_loss = vgg5(content_rgb)

        tF = vgg5(transfer) 
        art_loss, styleLoss, contentLoss = criterion(tF, sF_loss, cF_loss)
        summary_writer.add_scalar("train/art_loss", art_loss, global_step=iteration)
        summary_writer.add_scalar("train/styleLoss", styleLoss, global_step=iteration)
        summary_writer.add_scalar("train/contentLoss", contentLoss, global_step=iteration)

        tF_d = vgg5(transfer_d) 
        art_loss_d, styleLoss_d, contentLoss_d = criterion(tF_d, sF_loss, cF_loss)
        summary_writer.add_scalar("train/art_loss_d", art_loss_d, global_step=iteration)
        summary_writer.add_scalar("train/styleLoss_d", styleLoss_d, global_step=iteration)
        summary_writer.add_scalar("train/contentLoss_d", contentLoss_d, global_step=iteration)

        # tF_s = vgg5(transfer_s) 
        # art_loss_s, styleLoss_s, contentLoss_s = criterion(tF_s, sF_loss, cF_loss)
        # summary_writer.add_scalar("train/art_loss_s", art_loss_s, global_step=iteration)
        # summary_writer.add_scalar("train/styleLoss_s", styleLoss_s, global_step=iteration)
        # summary_writer.add_scalar("train/contentLoss_s", contentLoss_s, global_step=iteration)
        art_loss_s = 0.

        # 风格迁移loss
        total_art_loss = art_loss * 3.0 + art_loss_d + art_loss_s
        summary_writer.add_scalar("train/total_art_loss", total_art_loss, global_step=iteration)

        # 去噪loss
        tP = vgg5(propagated)
        contentV = content_rgb
        contentV_norm = contentV - torch.min(contentV)
        contentV_norm = contentV_norm / torch.max(contentV_norm)
        propagated_norm = propagated - torch.min(propagated)
        propagated_norm = propagated_norm / torch.max(propagated_norm)

        corr_loss, corr_img = criterion_corr(propagated_norm, contentV_norm) 
        corr_loss = 0.1 * corr_loss
        summary_writer.add_scalar("train/corr_loss", corr_loss, global_step=iteration)

        grad_loss, _, _ = criterion_grad(propagated_norm, contentV_norm)
        grad_loss = 0.5 * grad_loss
        summary_writer.add_scalar("train/grad_loss", grad_loss, global_step=iteration)

        contenProp_loss = criterion_contentProp(tP["r41"], cF_loss["r41"])

        # TV loss
        tv_image_loss = criterion_tv(recon_rgb)
        tv_image_loss_d = criterion_tv(recon_rgb_d)
        tv_image_loss_s = criterion_tv(recon_rgb_s)
        total_tv_image_loss = tv_image_loss * 3.0 + tv_image_loss_d * 1.0 + tv_image_loss_s * 1.0
        summary_writer.add_scalar("train/tv_image_loss", tv_image_loss, global_step=iteration)
        summary_writer.add_scalar("train/tv_image_loss_d", tv_image_loss_d, global_step=iteration)
        summary_writer.add_scalar("train/tv_image_loss_s", tv_image_loss_s, global_step=iteration)
        summary_writer.add_scalar("train/total_tv_image_loss", total_tv_image_loss, global_step=iteration)

        with torch.no_grad():

            transfer_pse_gt = []
            corrected_pse_gt = []
            for style_img_idx in range(style_img.shape[0]):
                transfer_pse_gt_temp, corrected_pse_gt_temp = linear_wct.style_transfer(rgbs_train, style_img[style_img_idx].unsqueeze(dim=0))
                transfer_pse_gt.append(transfer_pse_gt_temp)
                corrected_pse_gt.append(corrected_pse_gt_temp)
            transfer_pse_gt = torch.cat(transfer_pse_gt, dim=0)
            corrected_pse_gt = torch.cat(corrected_pse_gt, dim=0)
            # transfer_pse_gt, corrected_pse_gt = linear_wct.style_transfer(rgbs_train, style_img)

            # print(f"transfer_pse_gt: {transfer_pse_gt.min()} {transfer_pse_gt.max()}")
            # print(f"corrected_pse_gt: {corrected_pse_gt.min()} {corrected_pse_gt.max()}")
            # corrected_pse_gt = corrected_pse_gt.clamp(0,1)
        pse_transfer_loss = torch.mean((transfer - transfer_pse_gt) ** 2) * 200
        pse_transfer_loss_d = torch.mean((transfer_d - transfer_pse_gt) ** 2)  * 200
        # pse_reg_loss_s = torch.mean((transfer_s - transfer_pse_gt) ** 2)  * 100
        pse_corrected_loss = torch.mean((propagated - corrected_pse_gt) ** 2)  * 200
        total_pse_loss = pse_transfer_loss + pse_transfer_loss_d + pse_corrected_loss
        summary_writer.add_scalar("train/pse_transfer_loss", pse_transfer_loss, global_step=iteration)
        summary_writer.add_scalar("train/pse_transfer_loss_d", pse_transfer_loss_d, global_step=iteration)
        summary_writer.add_scalar("train/pse_corrected_loss", pse_corrected_loss, global_step=iteration)
        summary_writer.add_scalar("train/total_pse_loss", total_pse_loss, global_step=iteration)

        # Temporal loss
        transfer_w_gt, forward_flow, cF_w, forward_flow_feature = criterion_temporal.GenerateFakeFrameAndFeature(transfer, cF)
        feature_w, transmatrix_w = matrix(canonical_features, cF_w, sF)
        transfer_w = dec(feature_w)  # 1 x 3 x 256 x 256
        temporal_transfer_loss, _ = criterion_temporal(transfer, transfer_w, forward_flow)
        summary_writer.add_scalar("train/temporal_transfer_loss", temporal_transfer_loss, global_step=iteration)

        propagated_w_gt, forward_flow_2, cF_w_2, forward_flow_feature_2 = criterion_temporal.GenerateFakeFrameAndFeature(propagated, cF)
        feature_w_2, transmatrix_w_2 = matrix(canonical_features, cF_w_2, sF)
        transfer_w_2 = dec(feature_w)  # 1 x 3 x 256 x 256
        content_rgb_w = criterion_temporal.flow_warp(content_rgb, forward_flow_2)
        propagated_w = spn(content_rgb_w, transfer_w_2)
        temporal_propagated_loss, _ = criterion_temporal(propagated, propagated_w, forward_flow_2)
        summary_writer.add_scalar("train/temporal_propagated_loss", temporal_propagated_loss, global_step=iteration)

        total_temporal_loss = temporal_transfer_loss + temporal_propagated_loss
        summary_writer.add_scalar("train/total_temporal_loss", total_temporal_loss, global_step=iteration)


        total_loss = total_art_loss + contenProp_loss + corr_loss + grad_loss + total_tv_image_loss  + total_pse_loss + total_temporal_loss

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * lr_factor

        # Print the current values of the losses.
        if iteration % args.progress_refresh_rate == 0:
            pbar.set_description(
                f'Iteration {iteration:05d}:'
                # + f' psnr_p = {PSNR_pixel:.2f}'
                # + f' psnr_f = {PSNR_feature:.2f}'
            )

        if iteration % 100 == 0:
            summary_writer.add_image('recon_rgb', recon_rgb.clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('recon_rgb_d', recon_rgb_d.clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('recon_rgb_s', recon_rgb_s.clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('transfer', transfer[0].clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('transfer_d', transfer_d[0].clamp(0, 1).squeeze(), global_step=iteration)
            # summary_writer.add_image('transfer_s', transfer_s.clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('rgbs_train', rgbs_train.clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('style_img', style_img[0].clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('propagated', propagated[0].clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('transfer_pse_gt', transfer_pse_gt[0].clamp(0, 1).squeeze(), global_step=iteration)
            summary_writer.add_image('corrected_pse_gt', corrected_pse_gt[0].clamp(0, 1).squeeze(), global_step=iteration)

        if iteration % 100 == 0:
            matrix_path = f"{logfolder}/{args.expname}_matrix.th"
            spn_path = f"{logfolder}/{args.expname}_spn.th"
            torch.save(matrix.state_dict(), matrix_path)
            torch.save(spn.state_dict(), spn_path)
        # if iteration % (args.progress_refresh_rate*20) == 1:
        # if iteration % (args.progress_refresh_rate*2) == 1:
            # summary_writer.add_image('output', make_grid([rgbs_train.squeeze(), 
            #                                               recon_rgb_denorm.clamp(0, 1).squeeze()],  
            #                                               nrow=2, padding=0, normalize=False),
            #                                               global_step=iteration)
        

    matrix_path = f"{logfolder}/{args.expname}_matrix.th"
    spn_path = f"{logfolder}/{args.expname}_spn.th"
    torch.save(matrix.state_dict(), matrix_path)
    torch.save(spn.state_dict(), spn_path)
    tensorf.save(
        poses_mtx.detach().cpu(),
        focal_refine.detach().cpu(),
        f"{logfolder}/{args.expname}.th",
    )
    tensorf_static.save(
        poses_mtx.detach().cpu(),
        focal_refine.detach().cpu(),
        f"{logfolder}/{args.expname}_static.th",
    )



if __name__ == "__main__":
    torch.set_default_dtype(torch.float32)
    torch.manual_seed(20211202)
    np.random.seed(20211202)

    args = config_parser()
    print(args)

    # 限制cpu core使用数
    import os
    from multiprocessing import cpu_count

    # cpu_num = cpu_count() // 2 # 自动获取最大核心数目
    cpu_num = int(cpu_count() * args.cpu_percentage)
    os.environ ['OMP_NUM_THREADS'] = str(cpu_num)
    os.environ ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
    os.environ ['MKL_NUM_THREADS'] = str(cpu_num)
    os.environ ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
    os.environ ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
    torch.set_num_threads(cpu_num)

    if args.export_mesh:
        export_mesh(args)

    if args.render_only and (args.render_test or args.render_path):
        render_test(args, os.path.join(args.basedir, args.expname))
    else:
        reconstruction(args)