import os
import torch
from abc import ABC, abstractmethod
import wandb
import options
import utils
from pytorch_msssim import ssim, SSIM
import numpy as np
import torchvision
from tqdm import tqdm
import lpips
from imaginaire.losses import FeatureMatchingLoss, GaussianKLLoss, PerceptualLoss,GANLoss
import cv2
from imaginaire.utils.trainer import get_scheduler
from .geometry_transform import render_sat
from model import geometry_transform
import csv
class BaseModel(ABC):
"""This class is an abstract base class (ABC) for models.
To create a subclass, you need to implement the following five functions:
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
-- <set_input>: unpack data from dataset and apply preprocessing.
-- <forward>: produce intermediate results.
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
"""
def __init__(self, opt,wandb=None):
"""Initialize the BaseModel class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
When creating your custom class, you need to implement your own initialization.
In this function, you should first call <BaseModel.__init__(self, opt)>
Then, you need to define four lists:
-- self.loss_names (str list): specify the training losses that you want to plot and save.
-- self.model_names (str list): define networks used in our training.
"""
self.wandb = wandb
if opt.isTrain:
opt.save_dir =wandb.dir
options.save_options_file(opt,opt.save_dir)
self.opt = opt
self.device = "cpu" if opt.cpu or not torch.cuda.is_available() else "cuda:{}".format(opt.gpu)
# torch.backends.cudnn.benchmark = True
self.model_names = []
self.train_loader = None
self.val_loader = None
self.sty_loader = None
self.loss_fn_alex = lpips.LPIPS(net='alex',eval_mode=True).cuda()
if opt.task=='test':
self.loss_fn_sque = lpips.LPIPS(net='squeeze',eval_mode=True).cuda()
self.mseloss = torch.nn.MSELoss(True,True)
self.criteria = {}
self.weights = {}
if hasattr(opt.optim.loss_weight, 'GaussianKL'):
if opt.optim.loss_weight.GaussianKL:
self.criteria['GaussianKL'] = GaussianKLLoss()
self.weights['GaussianKL'] = opt.optim.loss_weight.GaussianKL
if hasattr(opt.optim.loss_weight, 'L1'):
if opt.optim.loss_weight.L1:
self.criteria['L1'] = torch.nn.L1Loss(True,True)
self.weights['L1'] = opt.optim.loss_weight.L1
if hasattr(opt.optim.loss_weight, 'L2'):
if opt.optim.loss_weight.L2:
self.criteria['L2'] = torch.nn.MSELoss(True,True)
self.weights['L2'] = opt.optim.loss_weight.L2
if hasattr(opt.optim.loss_weight, 'SSIM'):
if opt.optim.loss_weight.SSIM:
self.criteria['SSIM'] = SSIM(data_range =1., size_average=True, channel=3)
self.weights['SSIM'] = opt.optim.loss_weight.SSIM
if hasattr(opt.optim.loss_weight, 'Perceptual'):
if opt.optim.loss_weight.Perceptual:
self.criteria['Perceptual'] = \
PerceptualLoss(
network=opt.optim.perceptual_loss.mode,
layers=opt.optim.perceptual_loss.layers,
weights=opt.optim.perceptual_loss.weights).to(self.device)
self.weights['Perceptual'] = opt.optim.loss_weight.Perceptual
if hasattr(opt.optim.loss_weight, 'sky_inner'):
if opt.optim.loss_weight.sky_inner:
self.criteria['sky_inner'] = torch.nn.L1Loss(True,True)
self.weights['sky_inner'] = opt.optim.loss_weight.sky_inner
if hasattr(opt.optim.loss_weight, 'feature_matching'):
if opt.optim.loss_weight.feature_matching:
self.criteria['feature_matching'] = FeatureMatchingLoss()
self.weights['feature_matching'] = opt.optim.loss_weight.feature_matching
self.weights['GAN'] = opt.optim.loss_weight.GAN
self.criteria['GAN'] = GANLoss(gan_mode=opt.optim.gan_mode)
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new model-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
return parser
@abstractmethod
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): includes the data itself and its metadata information.
"""
pass
@abstractmethod
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
pass
def save_checkpoint(self,ep=0,latest=False):
"""
save trained models.
Args:
ep (int, optional): model epochs. Defaults to 0.
latest (bool, optional): qhether it is the latest model. Defaults to False.
"""
ckpt_save_path = os.path.join(self.wandb.dir,'checkpoint')
if not os.path.exists(ckpt_save_path):
os.mkdir(ckpt_save_path)
utils.save_checkpoint(self,ep=ep,latest=latest,output_path=ckpt_save_path)
if not latest:
print("checkpoint saved: {0}, epoch {1} ".format(self.opt.name,ep))
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
key = keys[i]
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
if module.__class__.__name__.startswith('InstanceNorm') and \
(key == 'running_mean' or key == 'running_var'):
if getattr(module, key) is None:
state_dict.pop('.'.join(keys))
if module.__class__.__name__.startswith('InstanceNorm') and \
(key == 'num_batches_tracked'):
state_dict.pop('.'.join(keys))
else:
self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
def setup_optimizer(self,opt):
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.optim.lr_gen, betas=(opt.optim.beta1, 0.999),eps=1.e-7)
if opt.isTrain:
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.optim.lr_dis, betas=(opt.optim.beta1, 0.999))
if opt.optim.lr_policy:
self.sch_G = get_scheduler(opt.optim, self.optimizer_G)
self.sch_D = get_scheduler(opt.optim, self.optimizer_D)
def optimize_parameters(self,opt):
self.netG.train()
# update Discriminators
self.backward_D(opt) # calculate gradients for D
# update Generator
self.backward_G(opt) # calculate graidents for G
psnr1 = -10*self.mseloss(self.fake_B.detach(),self.real_B.detach()).log10().item()
ssim_ = ssim(self.real_B.detach().float(), self.fake_B.detach().float(),data_range=1.)
out_dict = {
"train_ssim": ssim_,
"train_psnr1": psnr1,
}
# adjust learning rates according to schedule
if opt.optim.lr_policy:
out_dict["lr_D"]=self.sch_D.get_lr()[0]
out_dict["lr_G"]=self.sch_G.get_lr()[0]
out_dict.update(self.loss)
out_dict.update(self.dis_losses)
self.wandb.log(out_dict)
def validation(self,opt):
"""Used for validation and test in Center Ground-View Synthesis setting
Args:
opt (_type_): option dict
"""
print(10*"*","validate",10*"*")
self.netG.eval()
# six image reconstruction metrics
psnr_val = []
ssim_val = []
lpips_ale_val = []
lpips_squ_val = []
rmse_val = []
sd_val = []
with torch.no_grad():
# set the sky of all images with predefined sky histogram.
if opt.sty_img:
for _,data in enumerate(self.sty_loader):
self.set_input(data)
self.style_temp=self.sky_histc
break
for _,data in enumerate(tqdm(self.val_loader,ncols=100)):
self.set_input(data)
# if true: use the sky of predefined image
# if false: use the sky of corresponding GT
if opt.sty_img:
self.sky_histc = self.style_temp
self.forward(opt)
rmse = torch.sqrt(self.mseloss(self.fake_B*255.,self.real_B*255.)).item()
sd = sd_func(self.real_B,self.fake_B)
rmse_val.append(rmse)
sd_val.append(sd)
psnr1 = -10*self.mseloss(self.fake_B,self.real_B).log10().item()
ssim_ = ssim(self.real_B, self.fake_B,data_range=1.).item()
lpips_ale = torch.mean(self.loss_fn_alex((self.real_B*2.)-1, (2.*self.fake_B)-1)).cpu()
if opt.task=='test':
lpips_sque = torch.mean(self.loss_fn_sque((self.real_B*2.)-1, (2.*self.fake_B)-1)).cpu()
lpips_squ_val.append(lpips_sque)
psnr_val.append(psnr1)
ssim_val.append(ssim_)
lpips_ale_val.append(lpips_ale)
if opt.task in ['vis_test']:
if not os.path.exists(opt.vis_dir):
os.mkdir(opt.vis_dir)
sat_opacity,sat_depth = render_sat(opt,self.out_put['voxel'])
self.out_put['depth'] = (self.out_put['depth']/self.out_put['depth'].max())*255.
sat_depth = (sat_depth/sat_depth.max())*255.
for i in range(len(self.fake_B)):
depth_save = cv2.applyColorMap(self.out_put['depth'][i].squeeze().cpu().numpy().astype(np.uint8), cv2.COLORMAP_TURBO)
depth_sat_save = cv2.applyColorMap(sat_depth[i].squeeze().cpu().numpy().astype(np.uint8), cv2.COLORMAP_TURBO)
# cat generated ground images, GT ground images, predicted ground depth
torchvision.utils.save_image([self.fake_B[i].cpu(),self.real_B[i].cpu(),torch.flip(torch.from_numpy(depth_save).permute(2,0,1)/255.,[0])],os.path.join(opt.vis_dir,os.path.basename(self.image_paths[i])))
# cat GT satellite images, predicted satellite depth
torchvision.utils.save_image( [self.real_A[i].cpu() ,torch.flip(torch.from_numpy(depth_sat_save).permute(2,0,1)/255.,[0])],os.path.join(opt.vis_dir,os.path.basename(self.image_paths[i]).rsplit('.', 1)[0]+'_sat.jpg'))
# ground opacity
torchvision.utils.save_image([self.out_put['opacity'][i]] ,os.path.join(opt.vis_dir,os.path.basename(self.image_paths[i]).rsplit('.', 1)[0]+'_sat.jpg'))
psnr_avg = np.average(psnr_val)
ssim_avg = np.average(ssim_val)
lpips_ale_avg = np.average(lpips_ale_val)
if 'test' in opt.task:
lpips_squ_avg = np.average(lpips_squ_val)
rmse_avg = np.average(rmse_val)
sd_avg = np.average(sd_val)
if opt.task in ["train" , "Train"]:
out_dict = {
'val_psnr': psnr_avg,
'val_ssim': ssim_avg,
'val_lpips_ale':lpips_ale_avg,
'val_rmse':rmse_avg,
'val_sd':sd_avg
}
if opt.task=='test':
out_dict['val_lpips_squ'] = lpips_squ_avg
self.wandb.log(out_dict,commit=False)
else:
print(
{
'val_rmse':rmse_avg,
'val_ssim': ssim_avg,
'val_psnr': psnr_avg,
'val_sd':sd_avg,
'val_lpips_ale':lpips_ale_avg,
'val_lpips_squ':lpips_squ_avg,
}
)
with open('test_output.csv', mode='a', newline='') as csv_file:
writer = csv.writer(csv_file)
writer.writerow([rmse_avg, ssim_avg, psnr_avg, sd_avg, lpips_ale_avg, lpips_squ_avg])
def test_vid(self,opt):
"""Used for synthesis ground video
Args:
opt (_type_): option dict
"""
ckpt_list = os.listdir('wandb/')
for i in ckpt_list:
if opt.test_ckpt_path in i:
ckpt_path = i
ckpt = torch.load(os.path.join('wandb/',ckpt_path,'files/checkpoint/model.pth'))['netG']
print('load success!')
self.netG.load_state_dict(ckpt,strict=True)
self.netG.eval()
print(10*"*","test_video",10*"*")
pixels = []
if os.path.exists('vis_video/pixels.csv'):
with open('vis_video/pixels.csv', 'r') as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
x = float(row['x']) #x is
y = float(row['y'])
pixels.append((x, y))
else:
print('only render center point without vis_video/pixels.csv')
pixels = [(128,128)]
if opt.sty_img:
# inference with illumination from other images
for idx,data in enumerate(self.sty_loader):
self.set_input(data)
self.style_temp=self.sky_histc
break
with torch.no_grad():
for idx,data in enumerate(self.val_loader):
self.set_input(data)
if opt.sty_img:
self.sky_histc = self.style_temp
for i,(x,y) in enumerate(pixels):
opt.origin_H_W = [(y-128)/128 , (x-128)/128]
print(opt.origin_H_W)
self.forward(opt)
if not os.path.exists('vis_video'):
os.mkdir('vis_video')
# save voxel to visalize & satellite depth, works well on cvact
if i==0:
# pre-process for better visualize
volume_data = self.out_put.voxel.squeeze().cpu().numpy().transpose((1,2,0))
volume_data = np.clip(volume_data, None, 10)
import pyvista as pv
grid = pv.UniformGrid()
grid.dimensions = volume_data.shape
grid.spacing = (1, 1, 1)
grid.origin = (0, 0, 0)
grid.point_data['values'] = volume_data.flatten(order='F')
grid.save(os.path.join('vis_video',"volume_data.vtk") ) # vtk file could be visualized by ParaView app
sat_opacity,sat_depth = render_sat(opt,self.out_put['voxel'])
sat_depth = (2 - sat_depth)/(opt.data.max_height/15)*255.
depth_sat_save = cv2.applyColorMap(sat_depth[0].squeeze().cpu().numpy().astype(np.uint8), cv2.COLORMAP_TURBO)
torchvision.utils.save_image(torch.flip(torch.from_numpy(depth_sat_save).permute(2,0,1)/255.,[0]) ,os.path.join('vis_video',os.path.basename(self.image_paths[0])).replace('.png','_satdepth.png'))
torchvision.utils.save_image( [self.real_A[0].cpu() ] ,os.path.join('vis_video',os.path.basename(self.image_paths[0]).replace('.png','_sat.png')))
torchvision.utils.save_image( [self.real_B[0].cpu() ] ,os.path.join('vis_video',os.path.basename(self.image_paths[0]).replace('.png','_pano.png')))
self.out_put['depth'] = (self.out_put['depth']/self.out_put['depth'].max())*255.
depth_save = cv2.applyColorMap(self.out_put['depth'][0].squeeze().cpu().numpy().astype(np.uint8), cv2.COLORMAP_TURBO)
depth_save = torch.flip(torch.from_numpy(depth_save).permute(2,0,1)/255.,[0])
save_img = self.out_put.pred[0].cpu()
name = '%05d' % int(i) + ".png"
torchvision.utils.save_image(save_img,os.path.join('vis_video',os.path.basename(self.image_paths[0])).replace('.png',name))
save_img = depth_save
name = '%05d' % int(i) + "_depth.png"
torchvision.utils.save_image(save_img,os.path.join('vis_video',os.path.basename(self.image_paths[0])).replace('.png',name))
# save_img = self.out_put.generator_inputs[0][:3,:,:]
# name = '%05d' % int(i) + "_color_project.png"
# torchvision.utils.save_image(save_img,os.path.join('vis_video',os.path.basename(self.image_paths[0])).replace('.png',name))
def test_interpolation(self,opt):
"""Used for test interpolation
Args:
opt (_type_): option dict
"""
ckpt_list = os.listdir('wandb/')
for i in ckpt_list:
if opt.test_ckpt_path in i:
ckpt_path = i
ckpt = torch.load(os.path.join('wandb/',ckpt_path,'files/checkpoint/model.pth'))['netG']
print('load success!')
self.netG.load_state_dict(ckpt,strict=True)
self.netG.eval()
pixels = [(128,128)]
if opt.sty_img1:
for idx,data in enumerate(self.sty_loader1):
self.set_input(data)
self.style_temp1=self.sky_histc
break
if opt.sty_img2:
for idx,data in enumerate(self.sty_loader2):
self.set_input(data)
self.style_temp2=self.sky_histc
break
with torch.no_grad():
for idx,data in enumerate(self.val_loader):
self.set_input(data)
self.sky_histc1 = self.style_temp1
self.sky_histc2 = self.style_temp2
x,y = pixels[0]
opt.origin_H_W = [(y-128)/128 , (x-128)/128]
print(opt.origin_H_W)
estimated_height = self.netG.depth_model(self.real_A)
geo_outputs = geometry_transform.render(opt,self.real_A,estimated_height,self.netG.pano_direction,PE=self.netG.PE)
generator_inputs,opacity,depth = geo_outputs['rgb'],geo_outputs['opacity'],geo_outputs['depth']
if self.netG.gen_cfg.cat_opa:
generator_inputs = torch.cat((generator_inputs,opacity),dim=1)
if self.netG.gen_cfg.cat_depth:
generator_inputs = torch.cat((generator_inputs,depth),dim=1)
_, _, z1 = self.netG.style_encode(self.sky_histc1)
_, _, z2 = self.netG.style_encode(self.sky_histc2)
num_inter = 60
for i in range(num_inter):
z = z1 * (1-i/(num_inter-1)) + z2* (i/(num_inter-1))
z = self.netG.style_model(z)
output_RGB = self.netG.denoise_model(generator_inputs,z)
save_img = output_RGB.cpu()
name = 'img{:03d}.png'.format(i)
if not os.path.exists('vis_interpolation'):
os.mkdir('vis_interpolation')
torchvision.utils.save_image(save_img,os.path.join('vis_interpolation',name))
def test_speed(self,opt):
self.netG.eval()
random_input = torch.randn(1, 3, 256, 256).to(opt.device)
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
iterations = 300
times = torch.zeros(iterations)
with torch.no_grad():
for _ in range(50):
_ = self.netG(random_input,None,opt)
for iter in range(iterations):
starter.record()
_ = self.netG(random_input,None,opt)
ender.record()
torch.cuda.synchronize()
curr_time = starter.elapsed_time(ender) # 计算时间
times[iter] = curr_time
# print(curr_time)
mean_time = times.mean().item()
print("Inference time: {:.6f}, FPS: {} ".format(mean_time, 1000/mean_time))
def test_sty(self,opt):
ckpt_list = os.listdir('wandb/')
for i in ckpt_list:
if opt.test_ckpt_path in i:
ckpt_path = i
ckpt = torch.load(os.path.join('wandb/',ckpt_path,'files/checkpoint/model.pth'))['netG']
print('load success!')
self.netG.load_state_dict(ckpt,strict=True)
self.netG.eval()
print(10*"*","test_sty",10*"*")
self.netG.eval()
self.style_temp_list = []
with torch.no_grad():
num_val_loader = len(self.val_loader)
for i in range(num_val_loader):
for idx,data in enumerate(tqdm(self.val_loader,ncols=100)):
self.set_input(data)
if i==0:
self.style_temp_list.append(self.sky_histc)
name = '%05d' % int(idx)
torchvision.utils.save_image( [self.real_A[0].cpu() ] ,os.path.join(opt.vis_dir,os.path.basename(self.image_paths[0]).replace('.png',name+'_sat.png')))
self.sky_histc = self.style_temp_list[i]
self.forward(opt)
if not os.path.exists(opt.vis_dir):
os.mkdir(opt.vis_dir)
name = '%05d' % int(idx)+'_'+'%05d' % int(i)
name= name+ '.png'
torchvision.utils.save_image(self.fake_B[0].cpu(),os.path.join(opt.vis_dir, name))
def train(self,opt):
self.validation(opt)
for current_epoch in range(opt.max_epochs):
print(10*'-','current epoch is ',current_epoch,10*'-')
for idx,data in enumerate(tqdm(self.train_loader,ncols=100)):
self.set_input(data)
self.optimize_parameters(opt)
if idx%500==0 :
out_ing_dict = {
'train_input': wandb.Image(self.real_A[0].float()),
'train_pred_and_gt': wandb.Image(torch.cat([self.fake_B,self.real_B],2)[0].float()),
}
if hasattr(self.out_put, 'inter_RGB'):
out_ing_dict["train_inner_pred"] = wandb.Image(self.out_put.inter_RGB[0].float())
if opt.arch.gen.transform_mode in ['volum_rendering']:
out_ing_dict['train_inner_opacity'] = wandb.Image(self.out_put.opacity[0].float())
self.wandb.log(out_ing_dict,commit=False)
if opt.optim.lr_policy.iteration_mode:
self.sch_G.step()
self.sch_D.step()
if not opt.optim.lr_policy.iteration_mode:
self.sch_G.step()
self.sch_D.step()
self.validation(opt)
if current_epoch%5==0:
self.save_checkpoint(ep=current_epoch)
self.save_checkpoint(ep=current_epoch)
def test(self,opt):
ckpt_list = os.listdir('wandb/')
for i in ckpt_list:
if '.zip' not in i:
if opt.test_ckpt_path in i:
ckpt_path = i
ckpt = torch.load(os.path.join('wandb/',ckpt_path,'files/checkpoint/model.pth'))['netG']
print('load success!')
self.netG.load_state_dict(ckpt,strict=True)
# print(10*"*","validate",10*"*")
self.validation(opt)
print('if --task=vis_test,visible results will be saved,you can add "--vis_dir=xxx" to save in other dictionary',opt.vis_dir)
def _get_outputs(self, net_D_output, real=True):
r"""Return output values. Note that when the gan mode is relativistic.
It will do the difference before returning.
Args:
net_D_output (dict):
real_outputs (tensor): Real output values.
fake_outputs (tensor): Fake output values.
real (bool): Return real or fake.
"""
def _get_difference(a, b):
r"""Get difference between two lists of tensors or two tensors.
Args:
a: list of tensors or tensor
b: list of tensors or tensor
"""
out = list()
for x, y in zip(a, b):
if isinstance(x, list):
res = _get_difference(x, y)
else:
res = x - y
out.append(res)
return out
if real:
return net_D_output['real_outputs']
else:
return net_D_output['fake_outputs']
def sd_func(real, fake):
'''
ref: page 6 in https://arxiv.org/abs/1511.05440
'''
dgt1 = torch.abs(torch.diff(real,dim=-2))[:, :, 1:, 1:-1]
dgt2 = torch.abs(torch.diff(real, dim=-1))[:, :, 1:-1, 1:]
dpred1 = torch.abs(torch.diff(fake, dim=-2))[:, :, 1:, 1:-1]
dpred2 = torch.abs(torch.diff(fake, dim=-1))[:, :, 1:-1, 1:]
return 10*torch.log10(1.**2/torch.mean(torch.abs(dgt1+dgt2-dpred1-dpred2))).cpu().item()