spikezoo 0.1.2__py3-none-any.whl → 0.2__py3-none-any.whl

spikezoo/archs/stir/main.py

@@ -0,0 +1,397 @@
+import argparse
+import os
+import os.path as osp
+import shutil
+import time
+import numpy as np
+import torch
+import torch.optim
+import torch.backends.cudnn as cudnn
+from tensorboardX import SummaryWriter
+from thop import profile
+import pprint
+import datetime
+import lpips
+# import pyiqa
+# import cpbd
+import imageio
+from configs.yml_parser import *
+from datasets.dataset_sreds import *
+from configs.utils import *
+from metrics.psnr import *
+from metrics.ssim import *
+from metrics.losses import *
+from models.Vgg19 import *
+from spikingjelly.clock_driven import functional
+
+os.environ["KMP_BLOCKTIME"] = "0"
+os.environ["OMP_NUM_THREADS"] = "1"
+torch.set_num_threads(1)
+
+from models.networks_STIR import *
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--data_root', '-dr', type=str, default='/data/local_userdata/fanbin/REDS_dataset/REDS120fps')
+parser.add_argument('--arch', '-a', type=str, default='STIR')
+parser.add_argument('--batch_size', '-b', type=int, default=8)
+parser.add_argument('--learning_rate', '-lr', type=float, default=1e-4)
+parser.add_argument('--configs', '-cfg', type=str, default='./configs/STIR.yml')
+parser.add_argument('--epochs', '-ep', type=int, default=100)
+parser.add_argument('--epoch_size', '-es', type=int, default=1000)
+parser.add_argument('--workers', '-j', type=int, default=8)
+parser.add_argument('--pretrained', '-prt', type=str, default=None)
+parser.add_argument('--start_epoch', '-sep', type=int, default=0)
+parser.add_argument('--print_freq', '-pf', type=int, default=1)
+parser.add_argument('--save_dir', '-sd', type=str, default='ckpt_outputs')
+parser.add_argument('--save_name', '-sn', type=str, default='t1')
+parser.add_argument('--vis_path', '-vp', type=str, default='vis_train')
+parser.add_argument('--vis_name', '-vn', type=str, default='STIR_train')
+parser.add_argument('--eval_path', '-evp', type=str, default='vis_eval')
+parser.add_argument('--vis_freq', '-vf', type=int, default=200)
+parser.add_argument('--eval', '-e', action='store_true')
+parser.add_argument('--w_per', '-wper', type=float, default=0.2)
+parser.add_argument('--print_details', '-pd', action='store_true')
+parser.add_argument('--milestones', default=[20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], metavar='N', nargs='*')
+parser.add_argument('--lr_scale_factor', '-lrsf', type=float, default=0.7)
+parser.add_argument('--eval_interval', '-ei', type=int, default=5)
+parser.add_argument('--save_interval', '-si', type=int, default=5)
+parser.add_argument('--no_imwrite', action='store_true', default=False)
+args = parser.parse_args()
+
+args.milestones = [int(m) for m in args.milestones]
+print('milstones', args.milestones)
+
+cfg_parser = YAMLParser(args.configs)
+cfg = cfg_parser.config
+
+cfg['data']['root'] = args.data_root
+cfg = add_args_to_cfg(cfg, args, ['batch_size', 'arch', 'learning_rate', 'configs', 'epochs', 'epoch_size', 'workers', 'pretrained', 'start_epoch', 
+                        'print_freq', 'save_dir', 'save_name', 'vis_path', 'vis_name', 'eval_path', 'vis_freq', 'w_per'])
+
+n_iter = 0
+
+
+def train(cfg, train_loader, model, optimizer, epoch, train_writer):
+    ######################################################################
+    ## Init
+    global n_iter
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses_name = ['rec_loss', 'per_loss', 'mulscl_loss', 'all_loss']
+    losses = AverageMeter(precision=6, i=len(losses_name), names=losses_name)
+    model.train()
+    torch.cuda.synchronize()
+    end = time.time()
+    
+    vgg19 = Vgg19(requires_grad=False).cuda()
+    if torch.cuda.device_count() > 1:
+        vgg19 = nn.DataParallel(vgg19, list(range(torch.cuda.device_count())))
+
+    loss_fn_tv2 = VariationLoss(nc=2).cuda()
+    downsampleX2 = nn.AvgPool2d(2, stride=2).cuda()
+    loss_fn_L1 = L1Loss()
+    
+    ######################################################################
+    ## Training Loop
+    
+    for ww, data in enumerate(train_loader, 0):
+        
+        if ww >= args.epoch_size:
+            return
+
+        spikes = [spk.cuda() for spk in data['spikes']]
+        images = [img.cuda() for img in data['images']]
+        torch.cuda.synchronize()
+        data_time.update(time.time() - end)
+
+        cur_spks = torch.cat(spikes, dim=1)
+        
+        rec_loss = 0.0
+        per_loss = 0.0
+        loss_L1_multiscale = 0.0
+        loss_rep_est = 0.0
+
+        seq_len = len(data['spikes']) - 3###corres 23th img GT
+
+        for jj in range(1, 1+seq_len):
+            x = cur_spks[:, jj*20-11 : jj*20+50]
+
+            img_gt = images[jj+1]
+            
+            img_pred_0, Fs_lv_0, Fs_lv_1, Fs_lv_2, Fs_lv_3, Fs_lv_4, Est = model(x)
+            pred_F = [Fs_lv_0]
+            pred_F.append(Fs_lv_1)
+            pred_F.append(Fs_lv_2)
+            pred_F.append(Fs_lv_3)
+            pred_F.append(Fs_lv_4)
+            
+            # if jj > 1+2:
+            if jj >= 2:
+                rec_loss += loss_fn_L1(img_pred_0, img_gt, mean=True) / (seq_len - 1)
+                if cfg['train']['w_per'] > 0:
+                    per_loss += cfg['train']['w_per'] * compute_per_loss_single(img_pred_0, img_gt, vgg19) / (seq_len - 1)
+                else:
+                    per_loss = torch.tensor([0.0]).cuda()
+                
+                pyr_weights = [1.0, 0.5, 0.25, 0.25, 0.25] 
+                num=5  #pyramid: 3, 4, 5
+                for l in range(1, num):
+                    img_gt = downsampleX2(img_gt)
+                    loss_L1_multiscale += pyr_weights[l] * loss_fn_L1(pred_F[l][0], img_gt, mean=True) / (num-1) / (seq_len - 1)
+        all_loss = rec_loss + per_loss + loss_L1_multiscale #+ loss_rep_est
+        
+        # record loss
+        losses.update([rec_loss.item(), per_loss.item(), loss_L1_multiscale.item(), all_loss.item()])
+        train_writer.add_scalar('rec_loss', rec_loss.item(), n_iter)
+        train_writer.add_scalar('per_loss', per_loss.item(), n_iter)
+        train_writer.add_scalar('mulscl_loss', loss_L1_multiscale.item(), n_iter)
+        train_writer.add_scalar('total_loss', all_loss.item(), n_iter)
+
+        ## compute gradient and optimize
+        all_loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        functional.reset_net(model)
+        
+        torch.cuda.synchronize()
+        batch_time.update(time.time() - end)
+        torch.cuda.synchronize()
+        end = time.time()
+        n_iter += 1
+
+        if n_iter % cfg['train']['vis_freq'] == 0:
+            vis_img(cfg['train']['vis_path'], img_pred_0, cfg['train']['vis_name'])
+        
+        if ww % cfg['train']['print_freq'] == 0:
+            out_str = 'Epoch: [{:d}] [{:d}/{:d}],  Iter: {:d}  '.format(epoch, ww, len(train_loader), n_iter-1)
+            out_str += ' '.join(map('{:s} {:.4f} ({:.6f}) '.format, losses.names, losses.val, losses.avg))
+            out_str += 'lr {:.6f}'.format(optimizer.state_dict()['param_groups'][0]['lr'])
+            print(out_str)
+        
+        torch.cuda.synchronize()
+        end = time.time()
+    
+    return
+
+
+def validation(cfg, test_loader, model, epoch, auto_save_path):
+    global n_iter
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    metrics_name = ['PSNR', 'SSIM', 'LPIPS', 'AvgTime']
+    all_metrics = AverageMeter(i=len(metrics_name), precision=4, names=metrics_name)
+
+    timestamp1 = datetime.datetime.now().strftime('%m-%d')
+    timestamp2 = datetime.datetime.now().strftime('%H%M%S')
+
+    model.eval()
+
+    #lpips_loss = pyiqa.create_metric('lpips').cuda()
+    loss_fn_vgg = lpips.LPIPS(net='alex').cuda()
+    
+    padder = InputPadder(dims=(720, 1280))
+
+    for ww, data in enumerate(test_loader, 0):
+        torch.cuda.synchronize()
+        st1 = time.time()
+        spikes = torch.cat([spk.cuda() for spk in data['spikes']], dim=1)
+        images = data['images']
+        torch.cuda.synchronize()
+        data_time.update(time.time() - st1)
+
+        seq_metrics = AverageMeter(i=len(metrics_name), precision=4, names=metrics_name)
+
+        seq_len = len(data['spikes']) - 3###corres 23th img GT
+
+        pred_gif=[]
+        gt_gif=[]
+
+        for jj in range(1, 1+seq_len):
+            x = spikes[:, jj*20-11 : jj*20+50]
+            x = padder.pad(x)[0]
+            
+            gt = images[jj+1].cuda()
+
+            with torch.no_grad():
+                torch.cuda.synchronize()
+                st = time.time()
+                
+                out = model(x)
+                torch.cuda.synchronize()
+                mtime = time.time() - st
+            rec = padder.unpad(out)
+            
+            cur_rec = torch2numpy255(rec)
+            cur_gt = torch2numpy255(gt)
+
+            if not args.no_imwrite and args.eval:
+                save_path = osp.join(args.eval_path, timestamp1)
+                make_dir(save_path)
+                cur_vis_path = osp.join(save_path, '{:03d}_{:03d}.png'.format(ww, jj))
+                cv2.imwrite(cur_vis_path, cur_rec.astype(np.uint8))
+
+                pred_gif.append(cur_rec.astype(np.uint8))
+                gt_gif.append(cur_gt.astype(np.uint8))
+
+            cur_psnr = calculate_psnr(cur_rec, cur_gt)
+            cur_ssim = calculate_ssim(cur_rec, cur_gt)
+            with torch.no_grad():
+                cur_lpips = loss_fn_vgg(rec, gt)
+
+            cur_metrics_list = [cur_psnr, cur_ssim, cur_lpips.item(), mtime]
+            if args.eval:
+                print("[Seq%d, %d-th image]: PSNR:%.4f SSIM:%.4f LPIPS:%.4f Time:%.4f" % (ww, jj+2, cur_psnr, cur_ssim, cur_lpips.item(), mtime))
+
+            all_metrics.update(cur_metrics_list)
+            seq_metrics.update(cur_metrics_list)
+        
+        functional.reset_net(model)
+            
+        if args.print_details:
+            print('\n')
+            ostr = 'Data{:02d}  '.format(ww) + ' '.join(map('{:s} {:.4f} '.format, seq_metrics.names, seq_metrics.avg))
+            print(ostr)
+            print()
+    
+    ostr = 'All  ' + ' '.join(map('{:s} {:.4f} '.format, all_metrics.names, all_metrics.avg))
+    print(ostr)
+
+    if args.eval:
+        print('\n')
+    else:
+        print('Test current epoch\n')
+        f_metric_avg=open(os.path.join(auto_save_path, 'ckpt_'+args.save_name+'_metric_avg.txt'), 'a+')#Save the files next to the last line
+        f_metric_avg.write('%s  ' % (str(epoch).zfill(2)))
+        f_metric_avg.write(ostr)
+        f_metric_avg.write('\n')
+        f_metric_avg.close()
+
+    return
+    
+
+def main():
+    ##########################################################################################################
+    # Set random seeds
+    set_seeds(cfg['seed'])
+
+    # Create save path and logs
+    timestamp1 = datetime.datetime.now().strftime('%m-%d')
+    timestamp2 = datetime.datetime.now().strftime('%H%M%S')
+    if args.save_name == None:
+        save_folder_name = 'b{:d}_{:s}'.format(args.batch_size, timestamp2)
+    else:
+        save_folder_name = 'b{:d}_{:s}_{:s}'.format(args.batch_size, args.save_name, timestamp2)
+
+    save_path = osp.join(args.save_dir, timestamp1, save_folder_name)
+    print('save path: ', save_path)
+    if args.eval:
+        print('\n')
+    else:
+        make_dir(save_path)
+        #auto save test results during training
+        f_metric_avg=open(os.path.join(save_path, 'ckpt_'+args.save_name+'_metric_avg.txt'), 'w')
+        f_metric_avg.close()
+
+    make_dir(args.vis_path)
+    make_dir(args.eval_path)
+    
+    train_writer = SummaryWriter(save_path)
+
+    if args.eval:
+        shutil.rmtree(save_path)
+        print('remove path: ', save_path)
+
+    cfg_str = pprint.pformat(cfg)
+    print('=> configurations: ')
+    print(cfg_str)
+    
+    ##########################################################################################################
+    ## Create model
+    model = eval(args.arch)()
+
+    if args.pretrained:
+        network_data = torch.load(args.pretrained)
+        print('=> using pretrained model {:s}'.format(args.pretrained))
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+        model.load_state_dict(network_data)
+    else:
+        network_data = None
+        print('=> train from scratch')
+        model.init_weights()
+        print('=> model params: {:.6f}M'.format(model.num_parameters()/1e6))
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+
+    cudnn.benchmark = True
+
+    ##########################################################################################################
+    ## Create Optimizer
+    cfgopt = cfg['optimizer']
+    cfgmdl = cfg['model']
+    assert(cfgopt['solver'] in ['Adam', 'SGD'])
+    print('=> settings {:s} solver'.format(cfgopt['solver']))
+    
+    param_groups = [{'params': model.parameters(), 'weight_decay': cfgmdl['flow_weight_decay']}]
+    if cfgopt['solver'] == 'Adam':
+        optimizer = torch.optim.Adam(param_groups, args.learning_rate, betas=(cfgopt['momentum'], cfgopt['beta']))
+    elif cfgopt['solver'] == 'SGD':
+        optimizer = torch.optim.SGD(param_groups, args.learning_rate, momentum=cfgopt['momentum'])
+
+    ##########################################################################################################
+    ## Dataset
+    train_set = sreds_train(cfg)
+    train_loader = torch.utils.data.DataLoader(
+        train_set,
+        drop_last=False,
+        batch_size=cfg['train']['batch_size'],
+        shuffle=True,
+        num_workers=cfg['train']['workers'],
+        # pin_memory=True
+    )
+
+    test_set = sreds_test(cfg)
+    test_loader = torch.utils.data.DataLoader(
+        test_set,
+        drop_last=False,
+        batch_size=1,
+        shuffle=False,
+        num_workers=cfg['train']['workers']
+    )
+
+    ##########################################################################################################
+    ## Train or Evaluate
+    if args.eval:
+        validation(cfg=cfg, test_loader=test_loader, model=model, epoch=0, auto_save_path=save_path)
+    else:
+        epoch = cfg['train']['start_epoch']
+        while(True):
+            train(
+                cfg=cfg,
+                train_loader=train_loader,
+                model=model,
+                optimizer=optimizer,
+                epoch=epoch,
+                train_writer=train_writer
+            )
+            epoch += 1
+
+            # scheduler can be added here
+            if epoch in args.milestones:
+                for param_group in optimizer.param_groups:
+                    param_group['lr'] = param_group['lr'] * args.lr_scale_factor
+
+            # save model
+            if epoch % args.save_interval == 0:
+                model_save_name = '{:s}_epoch{:03d}.pth'.format(cfg['model']['arch'], epoch)
+                torch.save(model.state_dict(), osp.join(save_path, model_save_name))
+            
+            if epoch % args.eval_interval == 0:
+                validation(cfg=cfg, test_loader=test_loader, model=model, epoch=epoch, auto_save_path=save_path)
+
+            if epoch >= cfg['train']['epochs']:
+                break
+
+if __name__ == '__main__':
+    main()

spikezoo/archs/stir/metrics/losses.py

@@ -0,0 +1,219 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import torch.nn.functional as F
+
+import math
+
+from package_core.losses import *
+
+def compute_l1_loss(img_list, gt):
+    l1_loss = 0.0
+    for img in img_list:
+        cur_size = img.shape[-2:]
+        #gt_resize = F.interpolate(gt, size=cur_size, mode="bilinear", align_corners=False)
+        l1_loss += (img - gt).abs().mean()
+
+    return l1_loss
+
+def compute_per_loss_single(img, gt, vgg):
+    img_relu5_1 = vgg((img.repeat([1,3,1,1]) + 1.) / 2.)
+    with torch.no_grad():
+        gt_relu5_1 = vgg((gt.repeat([1,3,1,1]).detach() + 1.) / 2.)
+    percep_loss = F.mse_loss(img_relu5_1, gt_relu5_1)
+    return percep_loss
+
+def charbonier_loss(pred_im, im_gt, epsilon=0.001, mean=True):
+    x = pred_im - im_gt
+    loss = torch.mean(torch.sqrt(x ** 2 + epsilon ** 2))
+    return loss
+
+class L1Loss(nn.Module):
+    def __init__(self):
+        super(L1Loss, self).__init__()
+    def forward(self, output, target, weight=None, mean=False):
+        error = torch.abs(output - target)
+        if weight is not None:
+            error = error * weight.float()
+            if mean!=False:
+                return error.sum() / weight.float().sum()
+        if mean!=False:
+            return error.mean()
+        return error.sum()
+
+class VariationLoss(nn.Module):
+    def __init__(self, nc, grad_fn=Grid_gradient_central_diff):
+        super(VariationLoss, self).__init__()
+        self.grad_fn = grad_fn(nc)
+
+    def forward(self, image, weight=None, mean=False):
+        dx, dy = self.grad_fn(image)
+        variation = dx**2 + dy**2
+
+        if weight is not None:
+            variation = variation * weight.float()
+            if mean!=False:
+                return variation.sum() / weight.sum()
+        if mean!=False:
+            return variation.mean()
+        return variation.sum()
+
+class EdgeAwareVariationLoss(nn.Module):
+    def __init__(self, in1_nc, in2_nc, grad_fn=Grid_gradient_central_diff):
+        super(EdgeAwareVariationLoss, self).__init__()
+        self.in1_grad_fn = grad_fn(in1_nc)
+        self.in2_grad_fn = grad_fn(in2_nc)
+
+    def forward(self, in1, in2, mean=False):
+        in1_dx, in1_dy = self.in1_grad_fn(in1)
+        in2_dx, in2_dy = self.in2_grad_fn(in2)
+
+        abs_in1_dx, abs_in1_dy = in1_dx.abs().sum(dim=1,keepdim=True), in1_dy.abs().sum(dim=1,keepdim=True)
+        abs_in2_dx, abs_in2_dy = in2_dx.abs().sum(dim=1,keepdim=True), in2_dy.abs().sum(dim=1,keepdim=True)
+
+        weight_dx, weight_dy = torch.exp(-abs_in2_dx), torch.exp(-abs_in2_dy)
+
+        variation = weight_dx*abs_in1_dx + weight_dy*abs_in1_dy
+
+        if mean!=False:
+            return variation.mean()
+        return variation.sum()
+
+class PerceptualLoss():
+    def contentFunc(self):
+        conv_3_3_layer = 14
+        cnn = models.vgg19(pretrained=True).features
+        cnn = cnn.cuda()
+        model = nn.Sequential()
+        model = model.cuda()
+        for i,layer in enumerate(list(cnn)):
+            model.add_module(str(i),layer)
+            if i == conv_3_3_layer:
+                break
+        return model
+        
+    def __init__(self, loss):
+        self.criterion = loss
+        self.contentFunc = self.contentFunc()
+            
+    def get_loss(self, fakeIm, realIm):
+        f_fake = self.contentFunc.forward(fakeIm)
+        f_real = self.contentFunc.forward(realIm)
+        f_real_no_grad = f_real.detach()
+        loss = self.criterion(f_fake, f_real_no_grad)
+        return loss
+
+class SSIMLoss(nn.Module):
+    def __init__(self, nc=3):
+        super(SSIMLoss, self).__init__()
+        self.window_size=5
+        self.gaussian_img_kernel = self.create_gaussian_window(self.window_size, nc).float()
+
+    def create_gaussian_window(self, window_size, channel):
+        def _gaussian(window_size, sigma):
+            gauss = torch.Tensor(
+                [math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+            return gauss/gauss.sum()
+
+        _1D_window = _gaussian(window_size, 1.5).unsqueeze(1)
+        _2D_window = _1D_window@(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+        return window
+
+    def forward(self, img1, img2, mask=None):
+        self.gaussian_img_kernel = self.gaussian_img_kernel.to(img1.device)
+
+        params = {'weight': self.gaussian_img_kernel,
+              'groups': 3, 'padding': self.window_size//2}
+        mu1 = F.conv2d(img1, **params)
+        mu2 = F.conv2d(img2, **params)
+
+        mu1_sq = mu1.pow(2)
+        mu2_sq = mu2.pow(2)
+        mu1_mu2 = mu1*mu2
+
+        sigma1_sq = F.conv2d(img1*img1, **params) - mu1_sq
+        sigma2_sq = F.conv2d(img2*img2, **params) - mu2_sq
+        sigma12 = F.conv2d(img1*img2, **params) - mu1_mu2
+
+        C1 = 0.01**2
+        C2 = 0.03**2
+
+        ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2)) / ((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+        if mask is not None:
+            ssim_map = ssim_map * mask
+
+        return (1.-ssim_map.mean())*0.5
+
+def EPE3D_loss(input_flow, target_flow, mask=None):
+    """
+    :param the estimated optical / scene flow
+    :param the ground truth / target optical / scene flow
+    :param the mask, the mask has value 0 for all areas that are invalid
+    """
+    invalid = None
+    if mask is not None:
+        invalid = 1.-mask
+
+    epe_map = torch.norm(target_flow-input_flow,p=2,dim=1)
+    B = epe_map.shape[0]
+
+    invalid_flow = (target_flow != target_flow) # check Nan same as torch.isnan
+
+    mask = (invalid_flow[:,0,:,:] | invalid_flow[:,1,:,:] | invalid_flow[:,2,:,:]) 
+    if invalid is not None:
+        mask = mask | (invalid.view(mask.shape) > 0)
+
+    epes = []
+    for idx in range(B):
+        epe_sample = epe_map[idx][~mask[idx].data]
+        if len(epe_sample) == 0:
+            epes.append(torch.zeros(()).type_as(input_flow))
+        else:
+            epes.append(epe_sample.mean()) 
+
+    return torch.stack(epes)
+
+def compute_RT_EPE_loss(T_est, T_gt, depth0, K, mask=None): 
+    """ Compute the epe point error of rotation & translation
+    :param estimated rotation matrix Bx3x3
+    :param estimated translation vector Bx3
+    :param ground truth rotation matrix Bx3x3
+    :param ground truth translation vector Bx3
+    :param reference depth image, 
+    :param camera intrinsic 
+    """
+    R_est = T_est[:,:3,:3]
+    t_est = T_est[:,:3,3]
+    R_gt = T_gt[:,:3,:3]
+    t_gt = T_gt[:,:3,3]
+
+    loss = 0
+    if R_est.dim() > 3: # training time [batch, num_poses, rot_row, rot_col]
+        rH, rW = 60, 80 # we train the algorithm using a downsized input, (since the size of the input is not super important at training time)
+
+        B,C,H,W = depth0.shape
+        rdepth = func.interpolate(depth0, size=(rH, rW), mode='bilinear')
+        rmask = func.interpolate(mask.float(), size=(rH,rW), mode='bilinear')
+        rK = K.clone()
+        rK[:,0] *= float(rW) / W
+        rK[:,1] *= float(rH) / H
+        rK[:,2] *= float(rW) / W
+        rK[:,3] *= float(rH) / H
+        xyz = batch_inverse_project(rdepth, rK)
+        flow_gt = batch_transform_xyz(xyz, R_gt, t_gt, get_Jacobian=False)
+
+        for idx in range(R_est.shape[1]):
+            flow_est= batch_transform_xyz(xyz, R_est[:,idx], t_est[:,idx], get_Jacobian=False)
+            loss += EPE3D_loss(flow_est, flow_gt.detach(), rmask) #* (1<<idx) scaling does not help that much
+    else:
+        xyz = batch_inverse_project(depth0, K)
+        flow_gt = batch_transform_xyz(xyz, R_gt, t_gt, get_Jacobian=False)
+
+        flow_est= batch_transform_xyz(xyz, R_est, t_est, get_Jacobian=False)
+        loss = EPE3D_loss(flow_est, flow_gt, mask)
+
+    loss = loss.sum()/float(len(loss))
+    return loss

spikezoo/archs/stir/metrics/psnr.py

@@ -0,0 +1,22 @@
+import math
+import numpy as np
+
+# --------------------------------------------
+# PSNR
+# --------------------------------------------
+def calculate_psnr(img1, img2, border=0):
+    # img1 and img2 have range [0, 255]
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))

spikezoo/archs/stir/metrics/ssim.py

@@ -0,0 +1,54 @@
+import numpy as np
+import cv2
+
+# --------------------------------------------
+# SSIM
+# --------------------------------------------
+def calculate_ssim(img1, img2, border=0):
+    '''calculate SSIM
+    the same outputs as MATLAB's
+    img1, img2: [0, 255]
+    '''
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[2] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
+            return np.array(ssims).mean()
+        elif img1.shape[2] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+
+def ssim(img1, img2):
+    C1 = (0.01 * 255)**2
+    C2 = (0.03 * 255)**2
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()