视频超分、图像超分常用数据集Vimeo90K的下载及处理

最新推荐文章于 2024-10-06 06:00:00 发布

原创

最新推荐文章于 2024-10-06 06:00:00 发布 · 1.4w 阅读

24 ·

CC 4.0 BY-SA版权

文章标签：

#python

本文介绍了如何将Vimeo90K数据集转换为lmdb格式，包括GT和LR版本，以及提供通用图像处理工具和针对Vimeo90K、REDS等特定数据集的详细步骤。

vimeo90k数据集下载链接：

http://toflow.csail.mit.edu/

下载时点击最下方的训练集测试集一起下载即可，共计82G.数据集中已经划分了训练集和测试集，并给出了划分的txt以及4倍下采样后的图像。

在视频超分中，常常需要将图像转存为lmdb格式，该格式下运行速度更快，并且是依赖于缓存的。下面提供了将vimeo90K数据转换为lmdb格式的代码：create_lmdb.py

"""Create lmdb files for [General images (291 images/DIV2K) | Vimeo90K | REDS] training datasets"""

import sys
import os.path as osp
import glob
import pickle
from multiprocessing import Pool
import numpy as np
import lmdb
import cv2

sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
import data.util as data_util  # noqa: E402
import utils.util as util  # noqa: E402


def main():
    dataset = 'vimeo90k'  # vimeo90K | REDS | general (e.g., DIV2K, 291) | DIV2K_demo |test
    mode = 'GT'  # used for vimeo90k and REDS datasets
    # vimeo90k: GT | LR | flow
    # REDS: train_sharp, train_sharp_bicubic, train_blur_bicubic, train_blur, train_blur_comp
    #       train_sharp_flowx4
    if dataset == 'vimeo90k':
        vimeo90k(mode)
    elif dataset == 'REDS':
        REDS(mode)
    elif dataset == 'general':
        opt = {}
        opt['img_folder'] = '../../datasets/DIV2K/DIV2K800_sub'
        opt['lmdb_save_path'] = '../../datasets/DIV2K/DIV2K800_sub.lmdb'
        opt['name'] = 'DIV2K800_sub_GT'
        general_image_folder(opt)
    elif dataset == 'DIV2K':
        opt = {}
        ## GT
        opt['img_folder'] = '../../datasets/DIV2K_train_HR_sub'
        opt['lmdb_save_path'] = '../../datasets/DIV2K_train_HR_sub.lmdb'
        opt['name'] = 'DIV2K'
        general_image_folder(opt)
        ## LR
        opt['img_folder'] = '../../datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub'
        opt['lmdb_save_path'] = '../../datasets/DIV2K/imdb/DIV2K_train_LR_bicubic_X4_sub.lmdb'
        opt['name'] = 'DIV2Kx4'
        # general_image_folder(opt)
    elif dataset == 'test':
        test_lmdb('../../datasets/REDS/train_sharp_wval.lmdb', 'REDS')


def read_image_worker(path, key):
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    return (key, img)


def general_image_folder(opt):
    """Create lmdb for general image folders
    Users should define the keys, such as: '0321_s035' for DIV2K sub-images
    If all the images have the same resolution, it will only store one copy of resolution info.
        Otherwise, it will store every resolution info.
    """
    #### configurations
    read_all_imgs = False  # whether real all images to memory with multiprocessing
    # Set False for use limited memory
    BATCH = 5000  # After BATCH images, lmdb commits, if read_all_imgs = False
    n_thread = 40
    ########################################################
    img_folder = opt['img_folder']
    lmdb_save_path = opt['lmdb_save_path']
    meta_info = {'name': opt['name']}
    if not lmdb_save_path.endswith('.lmdb'):
        raise ValueError("lmdb_save_path must end with \'lmdb\'.")
    if osp.exists(lmdb_save_path):
        print('Folder [{:s}] already exists. Exit...'.format(lmdb_save_path))
        sys.exit(1)

    #### read all the image paths to a list
    print('Reading image path list ...')
    all_img_list = sorted(glob.glob(osp.join(img_folder, '*')))
    keys = []
    for img_path in all_img_list:
        keys.append(osp.splitext(osp.basename(img_path))[0])

    if read_all_imgs:
        #### read all images to memory (multiprocessing)
        dataset = {}  #