diff --git a/research/cv/LLVE-SEG/.keep b/research/cv/LLVE-SEG/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/research/cv/LLVE-SEG/requirements.txt b/research/cv/LLVE-SEG/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..4114ef2cf679744ec67a2bb241ade4b7ae3f4344 --- /dev/null +++ b/research/cv/LLVE-SEG/requirements.txt @@ -0,0 +1,8 @@ +argparse==1.1 +torch==1.2.0 +mindspore +torchvision==0.2.1 +numpy==1.18.1 +imageio==2.8.0 +opencv-python==4.2.0.32 +tensorboardX diff --git a/research/cv/LLVE-SEG/test.py b/research/cv/LLVE-SEG/test.py new file mode 100644 index 0000000000000000000000000000000000000000..b8aa04da509d6eb906e3b97d4a2377ba764a5d90 --- /dev/null +++ b/research/cv/LLVE-SEG/test.py @@ -0,0 +1,172 @@ +import os +import os.path as osp +import glob +import logging +import numpy as np +import cv2 +from PIL import Image + +import utils.util as util +import data.util as data_util +from models import create_model + +import os.path as osp +import logging +import time +import argparse +from collections import OrderedDict + +import options.options as option +import utils.util as util +from data.util import bgr2ycbcr +from data import create_dataset, create_dataloader +from models import create_model + +import time + +#### options +parser = argparse.ArgumentParser() +parser.add_argument('-opt', type=str, required=True, help='Path to options YMAL file.') +opt = option.parse(parser.parse_args().opt, is_train=False) +opt = option.dict_to_nonedict(opt) +import pdb + +def main(): + save_imgs = True + model = create_model(opt) + save_folder = './results_dynamic/{}'.format(opt['name']) + GT_folder = osp.join(save_folder, 'images/GT') + output_folder = osp.join(save_folder, 'images/output') + input_folder = osp.join(save_folder, 'images/input') + util.mkdirs(save_folder) + util.mkdirs(GT_folder) + util.mkdirs(output_folder) + util.mkdirs(input_folder) + + print('mkdir finish') + + util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) + logger = logging.getLogger('base') + + + for phase, dataset_opt in opt['datasets'].items(): + val_set = create_dataset(dataset_opt) + val_loader = create_dataloader(val_set, dataset_opt, opt, None) + + pbar = util.ProgressBar(len(val_loader)) + psnr_rlt = {} # with border and center frames + psnr_rlt_avg = {} + psnr_total_avg = 0. + + ssim_rlt = {} # with border and center frames + ssim_rlt_avg = {} + ssim_total_avg = 0. + + lpips_rlt = {} # with border and center frames + lpips_rlt_avg = {} + lpips_total_avg = 0. + + for val_data in val_loader: + folder = val_data['folder'][0] + idx_d = val_data['idx'] + if psnr_rlt.get(folder, None) is None: + psnr_rlt[folder] = [] + + if ssim_rlt.get(folder, None) is None: + ssim_rlt[folder] = [] + + if lpips_rlt.get(folder, None) is None: + lpips_rlt[folder] = [] + import time + start = time.time() + model.feed_data(val_data) + + model.test() + visuals = model.get_current_visuals() + end = time.time() + print(end-start) + pdb.set_trace() + rlt_img = util.tensor2img(visuals['rlt']) # uint8 + gt_img = util.tensor2img(visuals['GT']) # uint8 + + mid_ix = dataset_opt['N_frames'] // 2 + input_img = util.tensor2img(visuals['LQ']) #input_img = util.tensor2img(visuals['LQ'][mid_idx]) + #pdb.set_trace() + if save_imgs: + try: + tag = '{}.{}'.format(val_data['folder'], idx_d[0].replace('/', '-')) + print(osp.join(output_folder, '{}.png'.format(tag))) + cv2.imwrite(osp.join(output_folder, '{}.png'.format(tag)), rlt_img) + cv2.imwrite(osp.join(GT_folder, '{}.png'.format(tag)), gt_img) + + cv2.imwrite(osp.join(input_folder, '{}.png'.format(tag)), input_img) + + except Exception as e: + print(e) + import ipdb; ipdb.set_trace() + + # calculate PSNR + psnr = util.calculate_psnr(rlt_img, gt_img) + psnr_rlt[folder].append(psnr) + + ssim = util.calculate_ssim(rlt_img, gt_img) + #ssim = 0 + ssim_rlt[folder].append(ssim) + + lpips = util.calculate_lpips(rlt_img, gt_img) + lpips_rlt[folder].append(lpips) + + pbar.update('Test {} - {}'.format(folder, idx_d)) + for k, v in psnr_rlt.items(): + psnr_rlt_avg[k] = sum(v) / len(v) + psnr_total_avg += psnr_rlt_avg[k] + + for k, v in ssim_rlt.items(): + ssim_rlt_avg[k] = sum(v) / len(v) + ssim_total_avg += ssim_rlt_avg[k] + + for k, v in lpips_rlt.items(): + lpips_rlt_avg[k] = sum(v) / len(v) + lpips_total_avg += lpips_rlt_avg[k] + + psnr_total_avg /= len(psnr_rlt) + ssim_total_avg /= len(ssim_rlt) + lpips_total_avg /= len(lpips_rlt) + + log_s = '# Validation # PSNR: {:.4e}:'.format(psnr_total_avg) + for k, v in psnr_rlt_avg.items(): + log_s += ' {}: {:.4e}'.format(k, v) + logger.info(log_s) + + log_s = '# Validation # SSIM: {:.4e}:'.format(ssim_total_avg) + for k, v in ssim_rlt_avg.items(): + log_s += ' {}: {:.4e}'.format(k, v) + logger.info(log_s) + + psnr_all = 0 + psnr_count = 0 + for k, v in psnr_rlt.items(): + psnr_all += sum(v) + psnr_count += len(v) + psnr_all = psnr_all * 1.0 / psnr_count + print(psnr_all) + + ssim_all = 0 + ssim_count = 0 + for k, v in ssim_rlt.items(): + ssim_all += sum(v) + ssim_count += len(v) + ssim_all = ssim_all * 1.0 / ssim_count + print(ssim_all) + + lpips_all = 0 + lpips_count = 0 + for k, v in lpips_rlt.items(): + lpips_all += sum(v) + lpips_count += len(v) + lpips_all = lpips_all * 1.0 / lpips_count + print(lpips_all) + + +if __name__ == '__main__': + main() diff --git a/research/cv/LLVE-SEG/test_list.txt b/research/cv/LLVE-SEG/test_list.txt new file mode 100644 index 0000000000000000000000000000000000000000..302ede44188821e2fc82fb8965ce086cb5472605 --- /dev/null +++ b/research/cv/LLVE-SEG/test_list.txt @@ -0,0 +1,49 @@ +0013 +0009 +0020 +0007 +0010 +0012 +0154 +0036 +0037 +0039 +0047 +0048 +0049 +0050 +0051 +0052 +0065 +0066 +0075 +0076 +0078 +0088 +0092 +0099 +0091 +0103 +0104 +0105 +0106 +0107 +0151 +0145 +0147 +0139 +0129 +0083 +0153 +0157 +0180 +0175 +0181 +0196 +0170 +0166 +0172 +0177 +0167 +0169 +0191 diff --git a/research/cv/LLVE-SEG/train.py b/research/cv/LLVE-SEG/train.py new file mode 100644 index 0000000000000000000000000000000000000000..78c6c8f4cfa953d853eccde0df8685807b50623c --- /dev/null +++ b/research/cv/LLVE-SEG/train.py @@ -0,0 +1,189 @@ +import os +import math +import argparse +import random +import logging + +import mindspore.dataset as ds +import mindspore.dataset.transforms.c_transforms as C +import mindspore.dataset.vision.c_transforms as CV +import mindspore.nn as nn +import mindspore.ops as ops +import mindspore.common.dtype as mstype +from mindspore import context +from mindspore.communication.management import init, get_rank, get_group_size +from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor +from mindspore.train.model import Model +from mindspore.train.serialization import load_checkpoint, load_param_into_net + +import argparse +import math +import os +import random +import logging +import numpy as np +import time +import yaml + +from src.data import create_dataset, create_dataloader +from src.loss import GANLoss, GradientPenaltyLoss, CharbonnierLoss +from src.model import EDVR +from src.option import parse +import option +def init_dist(backend='nccl', **kwargs): + """initialization for distributed training""" + init(backend=backend, **kwargs) + +def main(): + #### options + parser = argparse.ArgumentParser() + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='./options/train/train_EDVR_ours.yml') + parser.add_argument('--launcher', choices=['none', 'mindspore'], default='none', help='job launcher') + parser.add_argument('--local_rank', type=int, default=0) + args = parser.parse_args() + + opt = parse(args.opt, is_train=True) + + #### distributed training settings + if args.launcher == 'none': # disabled distributed training + opt['dist'] = False + rank = -1 + print('Disabled distributed training.') + else: + opt['dist'] = True + init_dist() + rank = get_rank() + print('Enabled distributed training.') + + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + + #### loading resume state if exists + if opt['path'].get('resume_state', None): + # distributed resuming: all load into default GPU + device_id = rank % opt['num_gpus'] + resume_state = load_checkpoint(opt['path']['resume_state']) + option.check_resume(opt, resume_state['iter']) # check resume options + else: + resume_state = None + + #### mkdir and loggers + if rank <= 0: # normal training (rank -1) OR distributed training (rank 0) + if resume_state is None: + for key, path in opt['path'].items(): + if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key: + # config loggers. Before it, the log will not work + logger = logging.getLogger('base') + logger.setLevel(logging.INFO) + formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + file_handler = logging.FileHandler(os.path.join(opt['path']['log'], 'train_' + opt['name'] + '.log')) + file_handler.setLevel(logging.INFO) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + logger.info(option.dict2str(opt)) + # tensorboard logger + if opt['use_tb_logger'] and 'debug' not in opt['name']: + from mindspore.train.summary.summary_record import SummaryRecord + tb_logger = SummaryRecord(log_dir='./tb_logger/' + opt['name']) + else: + logger = logging.getLogger('base') + logger.setLevel(logging.INFO) + formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + #### random seed + seed = opt['train']['manual_seed'] + if seed is None: + seed = random.randint(1, 10000) + if rank <= 0: + logger.info('Random seed: {}'.format(seed)) + np.random.seed(seed) + random.seed(seed) + + #### create train and val dataloader + dataset_ratio = 200 # enlarge the size of each epoch + for phase, dataset_opt in opt['datasets'].items(): + if phase == 'train': + train_set = create_dataset(dataset_opt) + train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size'])) + total_iters = int(opt['train']['niter']) + total_epochs = int(math.ceil(total_iters / train_size)) + if opt['dist']: + train_sampler = ds.DistributedSampler(rank=get_rank(), num_shards=get_group_size(), shuffle=True) + train_sampler.set_num_replicas(get_group_size()) + train_sampler.set_seed(seed) + train_sampler.set_epoch(0) + train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler) + total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio))) + else: + train_loader = create_dataloader(train_set, dataset_opt, opt, None) + if rank <= 0: + logger.info('Number of train images: {:,d}, iters: {:,d}'.format(len(train_set), train_size)) + logger.info('Total epochs needed: {:d} for iters {:,d}'.format(total_epochs, total_iters)) + elif phase == 'val': + val_set = create_dataset(dataset_opt) + val_loader = create_dataloader(val_set, dataset_opt, opt, None) + if rank <= 0: + logger.info('Number of val images in [{:s}]: {:d}'.format(dataset_opt['name'], len(val_set))) + else: + raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase)) + assert train_loader is not None + + #### create model + model = EDVR(opt) + + #### resume training + if resume_state: + logger.info('Resuming training from epoch: {}, iter: {}.'.format(resume_state['epoch'], resume_state['iter'])) + start_epoch = resume_state['epoch'] + current_step = resume_state['iter'] + del resume_state + else: + current_step = 0 + start_epoch = 0 + + #### training + logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step)) + lr = ops.LearningRateSchedule(learning_rate=opt['train']['lr'], decay_steps=opt['train']['lr_decay'], decay_rate=opt['train']['lr_gamma'], staircase=True) + optimizer = nn.Adam(params=model.trainable_params(), learning_rate=lr, beta1=opt['train']['beta1'], beta2=opt['train']['beta2'], eps=opt['train']['epsilon']) + loss_fn = CharbonnierLoss() + model_with_loss = nn.WithLossCell(model, loss_fn) + model_with_loss = nn.TrainOneStepCell(model_with_loss, optimizer) + model = Model(model_with_loss) + + ckpt_config = CheckpointConfig(save_checkpoint_steps=opt['logger']['save_checkpoint_freq'], keep_checkpoint_max=opt['logger']['keep_checkpoint_max']) + ckpt_cb = ModelCheckpoint(prefix='EDVR', directory=opt['path']['models'], config=ckpt_config) + loss_cb = LossMonitor(per_print_times=opt['logger']['print_freq']) + + for epoch in range(start_epoch, total_epochs + 1): + if opt['dist']: + train_sampler.set_epoch(epoch) + for _, train_data in enumerate(train_loader): + current_step += 1 + if current_step > total_iters: + break + #### update learning rate + lr = ops.LearningRateSchedule(learning_rate=opt['train']['lr'], decay_steps=opt['train']['lr_decay'], decay_rate=opt['train']['lr_gamma'], staircase=True) + optimizer = nn.Adam(params=model.trainable_params(), learning_rate=lr, beta1=opt['train']['beta1'], beta2=opt['train']['beta2'], eps=opt['train']['epsilon']) + #### training + loss = model.train_step(train_data) + #### log + if current_step % opt['logger']['print_freq'] == 0: + logger.info('Epoch[{}], Step[{}], Loss {:.6f}'.format(epoch, current_step, loss.asnumpy())) + #### save models and training states + if current_step % opt['logger']['save_checkpoint_freq'] == 0: + ckpt_cb.save_checkpoint(epoch) + if rank <= 0: + logger.info('Saving the final model.') + ckpt_cb.save_checkpoint(total_epochs) + logger.info('End of training.') + tb_logger.close() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git "a/research/cv/LLVE-SEG/\350\257\264\346\230\216\346\226\207\346\241\243.docx" "b/research/cv/LLVE-SEG/\350\257\264\346\230\216\346\226\207\346\241\243.docx" new file mode 100644 index 0000000000000000000000000000000000000000..4c658e15591d98296ab19af349ee3880181bf022 Binary files /dev/null and "b/research/cv/LLVE-SEG/\350\257\264\346\230\216\346\226\207\346\241\243.docx" differ