The code for creating dataset

[coolKeen]

Hi, thanks for your nice work! I want to make my own dataset. Could you please share your code for creating the dataset?

[sooyekim]

Hi @coolKeen ,
Thanks for the interest in our work!
The Matlab code for creating the training dataset looks something like this:

clear all
%%%====== Data Setting ======%%%
format = '420';
scale = 2;
width = 1920; % LR
height = 1080; % LR
%%%====== Extraction Settings ======%%%
subim_size = 192;
depth = 5;
range = 9;
fr_stride = 10;
im_per_fr = 10;
%%%====== Directories ======%%%
dirname = 'your_YUV_data_directory_path';
file_HR(1).file = fullfile(dirname,'your_yuv_file.yuv');
file_HR(2).file = fullfile(dirname,'LG Slam Dunk 4K Demo.yuv');
scene_change = csvread('scene_change.csv');  % .csv file with scene change frame indices
%%%====== Initialization ======%%%
wid_crop = width-mod(width, subim_size);
hei_crop = height-mod(height, subim_size);
num = floor(hei_crop*wid_crop/subim_size/subim_size);

[fwidth,fheight] = yuv_factor(format);
framesLR = uint8(zeros(subim_size*subim_size, im_per_fr, depth));
LR_data = uint8(zeros(subim_size, subim_size, 3, depth, 1));
HR_data = uint8(zeros(subim_size*scale, subim_size*scale, 3, range-2, 1));

scene_change = scene_change + 1;
sc_valid(1).num = [2:3, 5:7, 9:28, 30:31, 33:36]; % Surfing
sc_valid(2).num = [2:15, 17:27, 30:32, 35:45, 47, 49, 51:59, 61:62, 65, 67, 69, 71]; % Slam Dunk

count = 1;
%%%====== Training Data Extraction ======%%%
for vid = 1:2
    for sc = sc_valid(vid).num
        for fr = scene_change(vid, sc)+ceil(range/2)-1 :fr_stride:scene_change(vid, sc + 1)-ceil(range/2)
            order = randperm(num,im_per_fr);
            for d = 1:range
                HR = loadYUV(width*scale,height*scale,fr-ceil(range/2)+d,file_HR(vid).file,fheight,fwidth, 'SDR');
                HR = modcrop(HR, subim_size*scale);
                if mod(d, 2) == 1 % LR
                    LR = imresize(HR, 1/scale);
                    TempLR(:, 1, :) = im2col(LR(:, :, 1), [subim_size subim_size], 'distinct');
                    TempLR(:, 2, :) = im2col(LR(:, :, 2), [subim_size subim_size], 'distinct');
                    TempLR(:, 3, :) = im2col(LR(:, :, 3), [subim_size subim_size], 'distinct');
                    framesLR = TempLR(:, :, order);
                    LR_data(:, :, :, ceil(d/2), count:count+im_per_fr-1)=reshape(framesLR, subim_size, subim_size, 3, im_per_fr);
                end
                if d > 1 && d < range % HR
                    TempHR(:, 1, :) = im2col(HR(:, :, 1), [subim_size*scale subim_size*scale], 'distinct');
                    TempHR(:, 2, :) = im2col(HR(:, :, 2), [subim_size*scale subim_size*scale], 'distinct');
                    TempHR(:, 3, :) = im2col(HR(:, :, 3), [subim_size*scale subim_size*scale], 'distinct');
                    framesHR = TempHR(:, :, order);
                    HR_data(:, :, :, d - 1, count:count+im_per_fr-1)=reshape(framesHR, subim_size*scale, subim_size*scale, 3, im_per_fr);
                end
            end
            fr
            count=count+im_per_fr;
        end
    end
end

%%%====== Save ======%%%
save('LR_data.mat','LR_data','-v7.3');
save('HR_data.mat','HR_data','-v7.3');

FYI, loadYUV() is a function that loads yuv frames and modcrop() is a function that crops the frame so that it's divisible by its second argument (eg. modcrop(HR, 2) crops array HR so that it's divisible by 2).

LMK if you have any other questions!