我自己生成的数据集在GitHub，包含了训练、测试、验证集：

GitHub - wolverinn/HEVC-CU-depths-dataset: A dataset that contains the Coding Unit image files and their corresponding depths for HEVC intra-prediction.

关于数据集的准备部分，还有一些更加细节的工作需要完成。首先就是训练集，验证集和测试集的划分，在数据集中使用的所有YUV序列如下：

由于存储空间的限制以及相邻帧之间的高度相似，所以在训练的时候并不会使用从一个YUV文件中抽取出的所有帧，而是会随机抽取一些帧，参考下面的公式来选取帧：

每个YUV的第2帧和第27帧都被抽取，之后的第n个帧就按照上面的公式抽取，f代表YUV文件的总帧数。代码：

def crop_image_to_ctu(video_number):
    frames = len(os.listdir("{}\\temp-frames".format(WORKSPACE_PATH))) # 当前视频一共有多少帧
    random_frames = [2,27]
    n = int((25+frames/40)//4)
    for i in range(frames//n):
        f_index = 27 + n*(i+1)
        if f_index > frames:
            break
        else:
            random_frames.append(f_index) # 随机抽取帧，有一个公式得出抽取的帧的编号
    for image_file in os.listdir("{}\\temp-frames".format(WORKSPACE_PATH)):
        frame_number = int(image_file.split('_')[2])-1 # ffmpeg生成帧编号是从1开始，这里减1将编号变成从0开始和ctu分割信息对应
        if frame_number in random_frames:
            img = Image.open(os.path.join("{}\\temp-frames".format(WORKSPACE_PATH),image_file))
            img_width, img_height = img.size
            ctu_number_per_img = math.ceil(img_width / 64) * math.ceil(img_height / 64)
            for ctu_number in range(ctu_number_per_img):
                img_row = ctu_number // math.ceil(img_width / 64)
                img_colonm = ctu_number % math.ceil(img_height / 64)
                start_pixel_x = img_colonm * 64
                start_pixel_y = img_row * 64
                cropped_img = img.crop((start_pixel_x, start_pixel_y, start_pixel_x + 64, start_pixel_y + 64)) # 依次对抽取到的帧进行裁剪
                cropped_img.save("{}\\v_{}_{}_{}_.jpg".format(IMG_PATH,video_number,str(frame_number),str(ctu_number)))
            img.close()
            dump_ctu_file(video_number, str(frame_number)) # 将当前帧的所有ctu分割信息保存到新的文件，只保存抽取的帧的信息
        os.remove(os.path.join("{}\\temp-frames".format(WORKSPACE_PATH),image_file)) # 裁剪过后的帧就删掉
    print("Total frames extracted from video_{} : {}".format(video_number,len(random_frames)))

在生成数据集的过程中，我发现对于CTU的分割信息，如果将YUV的所有帧的分割信息保存到txt文件，这样一个YUV就会产生几十兆甚至1080P的YUV会产生一百多兆的文本文件，这样会导致训练的时候每加载一个CTU分割信息都会花费一定的时间，影响训练的效率。所以首先想到的是不用保存所有的分割信息，只需要保存之前抽取出的帧的CTU划分信息就可以了，这样可以大大减少存储空间。其次，还可以对保存的数据结构进行优化，如果保存在文本文件中，每次的读取方式就只能从第一行开始遍历，直到找到对应的帧的CTU，效率也很低。所以我将需要保存的CTU划分信息保存到了Python的字典(dict)中，结构如下：

video_0 = {
    "frame_2":{
        "ctu_0":[...] # 16x16的划分信息
        "ctu_1":[...]
        .
        .
        .
        "ctu_103":[...]
    }
    "frame_27":{
        ...
    }
}

这样可以快速根据图片是第几帧，第几个CTU找到对应的划分信息。将这个字典使用Python自带的持久化库pickle保存到v_0.pkl，在使用的时候可以方便地进行读取。代码如下：

def dump_ctu_file(video_number,frame_number):
    # 将抽取到的帧的所有ctu分割信息保存到pickle：{"frame_number_1":{"ctu_number_1":[...];"ctu_number_2":[...]};"frame_number_2":...}
    frame_detected = 0
    ctu_number = "0"
    temp_ctu = []
    f_pkl = open("v_{}.pkl".format(video_number), 'rb')
    video_dict = pickle.load(f_pkl)
    f_pkl.close()
    video_dict[frame_number] = {}
    with open(CtuInfo_FILENAME,'r') as f:
        for i,line in enumerate(f):
            if frame_detected == 0:
                if "frame" in line:
                    current_frame = line.split(':')[1]
                    if int(frame_number) == int(current_frame):
                        frame_detected = 1
            elif "frame" in line:
                break
            elif "ctu" in line:
                temp_ctu = []
                ctu_number = int(line.split(':')[1])
                video_dict[frame_number][str(ctu_number)] = []
            else:
                line_depths = line.split(' ')
                for index in range(16):
                    temp_ctu.append(int(line_depths[index]))
                    video_dict[frame_number][str(ctu_number)] = temp_ctu
    f_pkl = open("v_{}.pkl".format(video_number), 'wb')
    pickle.dump(video_dict, f_pkl)
    f_pkl.close()

数据预处理是神经网络训练中的一个很重要也很花费时间的环节，必须考虑到神经网络所需要的格式化的输入，如何将图片与标签对应，以及读取所花费的时间的优化。根据这些去对原始数据使用一定的规则进行处理。只有数据预处理部分做好了，才能在训练部分专注于神经网络而不是一些琐碎的数据处理规则。

我自己生成的数据集GitHub地址：

HEVC-CU-depths-dataset