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

spikezoo/archs/wgse/train.py

@@ -0,0 +1,261 @@
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from pytorch_wavelets import DWT1DForward
+
+from transform import Compose, RandomCrop, RandomRotationFlip
+from dataset import DatasetREDS
+from dwtnets import Dwt1dResnetX_TCN
+from utils import calculate_psnr, calculate_ssim, mkdir
+
+parser = argparse.ArgumentParser(description='AAAI - WGSE - REDS')
+parser.add_argument('-c', '--cuda', type=str, default='1', help='select gpu card')
+parser.add_argument('-b', '--batch_size', type=int, default=16)
+parser.add_argument('-e', '--epoch', type=int, default=600)
+parser.add_argument('-w', '--wvl', type=str, default='db8', help='select wavelet base function')
+parser.add_argument('-j', '--jlevels', type=int, default=5)
+parser.add_argument('-k', '--kernel_size', type=int, default=3)
+parser.add_argument('-l', '--logpath', type=str, default='WGSE-Dwt1dNet')
+parser.add_argument('-r', '--resume_from', type=str, default=None)
+parser.add_argument('--dataroot', type=str, default=None)
+
+args = parser.parse_args()
+os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda)
+
+resume_folder = args.resume_from
+batch_size = args.batch_size
+learning_rate = 1e-4
+train_epoch = args.epoch
+dataroot = args.dataroot
+
+opt = 'adam'
+opt_param = "{\"beta1\":0.9,\"beta2\":0.99,\"weight_decay\":0}"
+
+random_seed = True
+manual_seed = 123
+
+scheduler = "MultiStepLR"
+scheduler_param = "{\"milestones\": [400, 600], \"gamma\": 0.2}"
+
+wvlname = args.wvl
+j = args.jlevels
+ks = args.kernel_size
+
+if_save_model = False
+eval_freq = 1
+checkpoints_folder = args.logpath + '-' + args.wvl + '-' + str(args.jlevels) + '-' + 'ks' + str(ks)
+
+
+def progress_bar_time(total_time):
+    hour = int(total_time) // 3600
+    minu = (int(total_time) % 3600) // 60
+    sec = int(total_time) % 60
+    return '%d:%02d:%02d' % (hour, minu, sec)
+
+def main():
+
+    global batch_size, learning_rate, random_seed, manual_seed, opt, opt_param, if_save_model, checkpoints_folder
+
+    mkdir(os.path.join('logs', checkpoints_folder))
+
+    if random_seed:
+        seed = np.random.randint(0, 10000)
+    else:
+        seed = manual_seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    opt_param_dict = json.loads(opt_param)
+    scheduler_param_dict = json.loads(scheduler_param)
+
+    cfg = {}
+    cfg['rootfolder'] = os.path.join(dataroot, 'train')
+    cfg['spikefolder'] = 'input'
+    cfg['imagefolder'] = 'gt'
+    cfg['H'] = 250
+    cfg['W'] = 400
+    cfg['C'] = 41
+    train_set = DatasetREDS(cfg,
+        transform=Compose(
+                [
+                    RandomCrop(128),
+                    RandomRotationFlip(0.0, 0.5, 0.5)
+                ]),
+    )
+
+    cfg = {}
+    cfg['rootfolder'] = os.path.join(dataroot, 'val')
+    cfg['spikefolder'] = 'input'
+    cfg['imagefolder'] = 'gt'
+    cfg['H'] = 250
+    cfg['W'] = 400
+    cfg['C'] = 41
+    test_set = DatasetREDS(cfg)
+
+    print('train_set len', train_set.__len__())
+    print('test_set len', test_set.__len__())
+
+    train_data_loader = torch.utils.data.DataLoader(
+            dataset=train_set,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=16,
+            drop_last=True)
+    test_data_loader = torch.utils.data.DataLoader(
+        dataset=test_set,
+        batch_size=1,
+        shuffle=True,
+        num_workers=1,
+        drop_last=False)
+
+    print(train_data_loader)
+    print(test_data_loader)
+
+    item0 = train_set[0]
+    s = item0['spikes']
+    s = s[None, :, 0:1, 0:1]
+    dwt = DWT1DForward(wave=wvlname, J=j)
+    B, T, H, W = s.shape
+    s_r = rearrange(s, 'b t h w -> b h w t')
+    s_r = rearrange(s_r, 'b h w t -> (b h w) 1 t')
+    yl, yh = dwt(s_r)
+    yl_size = yl.shape[-1]
+    yh_size = [yhi.shape[-1] for yhi in yh]
+    
+    model = Dwt1dResnetX_TCN(inc=41, wvlname=wvlname, J=j, yl_size=yl_size, yh_size=yh_size, num_residual_blocks=3, norm=None, ks=ks)
+
+    
+    if args.resume_from:
+        print("loading model weights from ", resume_folder)
+        saved_state_dict = torch.load(os.path.join(resume_folder, 'model_best.pt'))
+        model.load_state_dict(saved_state_dict.module.state_dict())
+        print("Weighted loaded.")
+
+    model = torch.nn.DataParallel(model).cuda()
+
+    # optimizer
+    if opt.lower() == 'adam':
+        assert ('beta1' in opt_param_dict.keys() and 'beta2' in opt_param_dict.keys() and 'weight_decay' in opt_param_dict.keys())
+        betas = (opt_param_dict['beta1'], opt_param_dict['beta2'])
+        del opt_param_dict['beta1']
+        del opt_param_dict['beta2']
+        opt_param_dict['betas'] = betas
+        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, **opt_param_dict)
+    elif opt.lower() == 'sgd':
+        assert ('momentum' in opt_param_dict.keys() and 'weight_decay' in opt_param_dict.keys())
+        optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, **opt_param_dict)
+    else:
+        raise ValueError()
+
+    lr_scheduler = getattr(torch.optim.lr_scheduler, scheduler)(optimizer, **scheduler_param_dict)
+    best_psnr, best_ssim = 0.0, 0.0
+
+    for epoch in range(train_epoch+1):
+        print('Epoch %d/%d ... ' % (epoch, train_epoch))
+
+        model.train()
+        total_time = 0
+        f = open(os.path.join('logs', checkpoints_folder, 'log.txt'), "a")
+        for i, item in enumerate(train_data_loader):
+
+            start_time = time.time()
+
+            spikes = item['spikes'].cuda()
+            image = item['image'].cuda()
+            optimizer.zero_grad()
+            
+            pred = model(spikes)
+            
+            loss = F.l1_loss(image, pred)
+            loss.backward()
+            optimizer.step()
+
+            elapse_time = time.time() - start_time  
+            total_time += elapse_time
+
+            lr_list = lr_scheduler.get_last_lr()
+            lr_str = ""
+            for ilr in lr_list:
+                lr_str += str(ilr) + ' '
+            print('\r[training] %3.2f%% | %6d/%6d [%s<%s, %.2fs/it] | LOSS: %.4f | LR: %s' % (
+                float(i + 1) / int(len(train_data_loader)) * 100, i + 1, int(len(train_data_loader)),
+                progress_bar_time(total_time),
+                progress_bar_time(total_time / (i + 1) * int(len(train_data_loader))),
+                total_time / (i + 1),
+                loss.item(),
+                lr_str), end='')
+            f.write('[training] %3.2f%% | %6d/%6d [%s<%s, %.2fs/it] | LOSS: %.4f | LR: %s\n' % (
+                float(i + 1) / int(len(train_data_loader)) * 100, i + 1, int(len(train_data_loader)),
+                progress_bar_time(total_time),
+                progress_bar_time(total_time / (i + 1) * int(len(train_data_loader))),
+                total_time / (i + 1),
+                loss.item(),
+                lr_str))
+        
+        lr_scheduler.step()
+
+        print('')
+        if epoch % eval_freq == 0:
+            model.eval()
+            with torch.no_grad():
+                sum_ssim = 0.0
+                sum_psnr = 0.0
+                sum_num = 0
+                total_time = 0
+                for i, item in enumerate(test_data_loader):
+                    start_time = time.time()
+
+                    spikes = item['spikes'][:, 130:171, :, :].cuda()
+                    image = item['image'].cuda()
+
+                    pred = model(spikes)
+
+                    prediction = pred[0].permute(1,2,0).cpu().numpy()
+                    gt = image[0].permute(1,2,0).cpu().numpy()
+
+                    sum_ssim += calculate_ssim(gt * 255.0, prediction * 255.0)
+                    sum_psnr += calculate_psnr(gt * 255.0, prediction * 255.0)
+                    sum_num += 1
+                    elapse_time = time.time() - start_time
+                    total_time += elapse_time
+                    print('\r[evaluating] %3.2f%% | %6d/%6d [%s<%s, %.2fs/it]' % (
+                        float(i + 1) / int(len(test_data_loader)) * 100, i + 1, int(len(test_data_loader)),
+                        progress_bar_time(total_time),
+                        progress_bar_time(total_time / (i + 1) * int(len(test_data_loader))),
+                        total_time / (i + 1)), end='')
+                    f.write('[evaluating] %3.2f%% | %6d/%6d [%s<%s, %.2fs/it]\n' % (
+                    float(i + 1) / int(len(test_data_loader)) * 100, i + 1, int(len(test_data_loader)),
+                    progress_bar_time(total_time),
+                    progress_bar_time(total_time / (i + 1) * int(len(test_data_loader))),
+                    total_time / (i + 1)))
+                        
+                sum_psnr /= sum_num
+                sum_ssim /= sum_num
+
+            print('')
+            print('\r[Evaluation Result] PSNR: %.3f | SSIM: %.3f' % (sum_psnr, sum_ssim))
+            f.write('[Evaluation Result] PSNR: %.3f | SSIM: %.3f\n' % (sum_psnr, sum_ssim))
+
+        if if_save_model and epoch % eval_freq == 0:
+            print('saving net...')
+            torch.save(model, os.path.join('logs', checkpoints_folder) + '/model_epoch%d.pt' % epoch)
+            print('saved')
+
+        if sum_psnr > best_psnr or sum_ssim > best_ssim:
+            best_psnr = sum_psnr
+            best_ssim = sum_ssim
+            print('saving best net...')
+            torch.save(model, os.path.join('logs', checkpoints_folder) + '/model_best.pt')
+            print('saved')
+
+        f.close()
+
+
+if __name__ == '__main__':
+    main()

spikezoo/archs/wgse/transform.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn.functional as F
+from math import sin, cos, pi
+import numbers
+import random
+
+
+class Compose(object):
+    """Composes several transforms together.
+    Args:
+        transforms (list of ``Transform`` objects): list of transforms to compose.
+    Example:
+        >>> transforms.Compose([
+        >>>     transforms.CenterCrop(10),
+        >>>     transforms.ToTensor(),
+        >>> ])
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, x, y):
+        for t in self.transforms:
+            x, y = t(x, y)
+        return x, y
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += '\n'
+            format_string += '    {0}'.format(t)
+        format_string += '\n)'
+        return format_string
+
+
+class RandomCrop(object):
+    """Crop the tensor at a random location.
+    """
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    @staticmethod
+    def get_params(x, output_size):
+        w, h = x.shape[2], x.shape[1]
+        th, tw = output_size
+        assert(th <= h)
+        assert(tw <= w)
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+
+        return i, j, th, tw
+
+    def __call__(self, x, y):
+        """
+            x: [C x H x W] Tensor to be rotated.
+        Returns:
+            Tensor: Cropped tensor.
+        """
+        i, j, h, w = self.get_params(x, self.size)
+
+        return x[:, i:i + h, j:j + w], y[:, i:i + h, j:j + w]
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(size={0})'.format(self.size)
+
+
+class RandomRotationFlip(object):
+    """Rotate the image by angle.
+    """
+
+    def __init__(self, degrees, p_hflip=0.5, p_vflip=0.5):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            if len(degrees) != 2:
+                raise ValueError("If degrees is a sequence, it must be of len 2.")
+            self.degrees = degrees
+
+        self.p_hflip = p_hflip
+        self.p_vflip = p_vflip
+
+    @staticmethod
+    def get_params(degrees, p_hflip, p_vflip):
+        """Get parameters for ``rotate`` for a random rotation.
+        Returns:
+            sequence: params to be passed to ``rotate`` for random rotation.
+        """
+        angle = random.uniform(degrees[0], degrees[1])
+        angle_rad = angle * pi / 180.0
+
+        M_original_transformed = torch.FloatTensor([[cos(angle_rad), -sin(angle_rad), 0],
+                                                    [sin(angle_rad), cos(angle_rad), 0],
+                                                    [0, 0, 1]])
+
+        if random.random() < p_hflip:
+            M_original_transformed[:, 0] *= -1
+
+        if random.random() < p_vflip:
+            M_original_transformed[:, 1] *= -1
+
+        M_transformed_original = torch.inverse(M_original_transformed)
+
+        M_original_transformed = M_original_transformed[:2, :].unsqueeze(dim=0)  # 3 x 3 -> N x 2 x 3
+        M_transformed_original = M_transformed_original[:2, :].unsqueeze(dim=0)
+
+        return M_original_transformed, M_transformed_original
+
+    def __call__(self, x, y):
+        """
+            x: [C x H x W] Tensor to be rotated.
+        Returns:
+            Tensor: Rotated tensor.
+        """
+        assert(len(x.shape) == 3)
+
+        M_original_transformed, M_transformed_original = self.get_params(self.degrees, self.p_hflip, self.p_vflip)
+        affine_gridx = F.affine_grid(M_original_transformed, x.unsqueeze(dim=0).shape, align_corners=False)
+        transformedx = F.grid_sample(x.unsqueeze(dim=0), affine_gridx, align_corners=False)
+
+        affine_gridy = F.affine_grid(M_original_transformed, y.unsqueeze(dim=0).shape, align_corners=False)
+        transformedy = F.grid_sample(y.unsqueeze(dim=0), affine_gridy, align_corners=False)
+
+        return transformedx.squeeze(dim=0), transformedy.squeeze(dim=0)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
+        format_string += ', p_flip={:.2f}'.format(self.p_hflip)
+        format_string += ', p_vlip={:.2f}'.format(self.p_vflip)
+        format_string += ')'
+        return format_string

spikezoo/archs/wgse/utils.py

@@ -0,0 +1,128 @@
+import os
+import math
+import numpy as np
+import cv2
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+def RawToSpike(video_seq, h, w, flipud=True):
+
+    video_seq = np.array(video_seq).astype(np.uint8)
+    img_size = h*w
+    img_num = len(video_seq)//(img_size//8)
+    SpikeMatrix = np.zeros([img_num, h, w], np.uint8)
+    pix_id = np.arange(0,h*w)
+    pix_id = np.reshape(pix_id, (h, w))
+    comparator = np.left_shift(1, np.mod(pix_id, 8))
+    byte_id = pix_id // 8
+
+    for img_id in np.arange(img_num):
+        id_start = int(img_id)*int(img_size)//8
+        id_end = int(id_start) + int(img_size)//8
+        cur_info = video_seq[id_start:id_end]
+        data = cur_info[byte_id]
+        result = np.bitwise_and(data, comparator)
+        if flipud:
+            SpikeMatrix[img_id, :, :] = np.flipud((result == comparator))
+        else:
+            SpikeMatrix[img_id, :, :] = (result == comparator)
+
+    return SpikeMatrix
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+# https://github.com/twhui/SRGAN-pyTorch
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def mkdirs(paths):
+    if isinstance(paths, str):
+        mkdir(paths)
+    else:
+        for path in paths:
+            mkdir(path)
+
+
+# --------------------------------------------
+# 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))
+
+
+# --------------------------------------------
+# 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()
+

spikezoo/datasets/base_dataset.py

@@ -6,7 +6,7 @@ import numpy as np
 from spikezoo.utils.spike_utils import load_vidar_dat
 import re
 from dataclasses import dataclass, replace
-from typing import Literal
+from typing import Literal,Union
 import warnings
 import torch
 from tqdm import tqdm
@@ -19,7 +19,7 @@ class BaseDatasetConfig:
     "Dataset name."
     dataset_name: str = "base"
     "Directory specifying location of data."
-    root_dir: Path = Path(__file__).parent.parent / Path("data/base")
+    root_dir: Union[str,Path] = Path(__file__).parent.parent / Path("data/base")
     "Image width."
     width: int = 400
     "Image height."
@@ -108,7 +108,6 @@ class BaseDataset(Dataset):
             spike_name,
             height=self.cfg.height,
             width=self.cfg.width,
-            out_type="float",
             out_format="tensor",
         )
         return spike

spikezoo/metrics/__init__.py

@@ -11,7 +11,7 @@ import torch.nn.functional as F
 
 # todo with the union type
 metric_pair_names = ["psnr", "ssim", "lpips", "mse"]
-metric_single_names = ["niqe", "brisque", "piqe"]
+metric_single_names = ["niqe", "brisque", "piqe", "liqe_mix", "clipiqa"]
 metric_all_names = metric_pair_names + metric_single_names
 
 metric_single_list = {}

spikezoo/models/base_model.py

@@ -45,7 +45,6 @@ class BaseModel(nn.Module):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.net = self.build_network().to(self.device)
         self.net = nn.DataParallel(self.net) if cfg.multi_gpu == True else self.net
-        self.spike_size = None
         self.model_half_win_length: int = cfg.model_win_length // 2
 
     # ! Might lead to low speed training on the BSF.
@@ -104,8 +103,7 @@ class BaseModel(nn.Module):
                 :,
                 spike_mid - self.model_half_win_length : spike_mid + self.model_half_win_length + 1,
             ]
-        if self.spike_size == None:
-            self.spike_size = (spike.shape[2], spike.shape[3])
+        self.spike_size = (spike.shape[2], spike.shape[3])
         return spike
 
     def preprocess_spike(self, spike):

spikezoo/pipeline/base_pipeline.py

@@ -67,7 +67,7 @@ class Pipeline:
         """Pipeline setup."""
         # save folder
         self.thistime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")[:23]
-        self.save_folder = Path(__file__).parent.parent / Path(f"results") if len(self.cfg.save_folder) == 0 else self.cfg.save_folder
+        self.save_folder = Path(f"results") if len(self.cfg.save_folder) == 0 else self.cfg.save_folder
         mode_name = "train" if self.cfg._mode == "train_mode" else "detect"
         self.save_folder = (
             self.save_folder / Path(f"{mode_name}/{self.thistime}")
@@ -93,6 +93,7 @@ class Pipeline:
     def spk2img_from_dataset(self, idx=0):
         """Func---Save the recoverd image and calculate the metric from the given dataset."""
         # save folder
+        self.logger.info("*********************** spk2img_from_dataset ***********************")
         save_folder = self.save_folder / Path(f"spk2img_from_dataset/{self.dataset.cfg.dataset_name}_dataset/{self.dataset.cfg.split}/{idx:06d}")
         os.makedirs(str(save_folder), exist_ok=True)
 
@@ -106,9 +107,10 @@ class Pipeline:
             img = None
         return self._spk2img(spike, img, save_folder)
 
-    def spk2img_from_file(self, file_path, height, width, img_path=None, remove_head=False):
+    def spk2img_from_file(self, file_path, height = -1, width  = -1, img_path=None, remove_head=False):
         """Func---Save the recoverd image and calculate the metric from the given input file."""
         # save folder
+        self.logger.info("*********************** spk2img_from_file ***********************")
         save_folder = self.save_folder / Path(f"spk2img_from_file/{os.path.basename(file_path)}")
         os.makedirs(str(save_folder), exist_ok=True)
 
@@ -135,6 +137,7 @@ class Pipeline:
     def spk2img_from_spk(self, spike, img=None):
         """Func---Save the recoverd image and calculate the metric from the given spike stream."""
         # save folder
+        self.logger.info("*********************** spk2img_from_spk ***********************")
         save_folder = self.save_folder / Path(f"spk2img_from_spk/{self.thistime}")
         os.makedirs(str(save_folder), exist_ok=True)
 
@@ -188,7 +191,7 @@ class Pipeline:
         if self.cfg.save_metric == True:
             self.logger.info(f"----------------------Method: {model_name.upper()}----------------------")
             # paired metric
-            for metric_name in metric_all_names:
+            for metric_name in self.cfg.metric_names:
                 if img is not None and metric_name in metric_pair_names:
                     self.logger.info(f"{metric_name.upper()}: {cal_metric_pair(recon_img,img,metric_name)}")
                 elif metric_name in metric_single_names:
@@ -203,8 +206,7 @@ class Pipeline:
             if img is not None:
                 img = tensor2npy(img[0, 0])
                 cv2.imwrite(f"{save_folder}/sharp_img.png", img)
-        self.logger.info(f"Images are saved on the {save_folder}")
-
+            self.logger.info(f"Images are saved on the {save_folder}")
         return recon_img_copy
 
     def _post_process_img(self, model_name, recon_img, gt_img):

spikezoo/pipeline/train_pipeline.py

@@ -66,7 +66,7 @@ class TrainPipeline(Pipeline):
         save_folder = self.save_folder / Path("imgs") / Path(f"{epoch:06d}")
         os.makedirs(save_folder, exist_ok=True)
         for batch_idx, batch in enumerate(tqdm(self.dataloader)):
-            if batch_idx % (len(self.dataloader) // 4) == 0:
+            if batch_idx % (len(self.dataloader) // 4) != 0:
                 continue 
             batch = self.model.feed_to_device(batch)
             outputs = self.model.get_outputs_dict(batch)