HEVC自带的编码器在决定CTU的最佳分割深度的时候会花费很多时间，我们的目标是训练卷积神经网络来根据输入的帧预测这一帧图像上面的所有CTU的最佳分割深度，这样会比HEVC自带的编码器快很多。
我们需要做的主要有：
- 构建用于训练神经网络的数据集
- 数据预处理
- 构建卷积神经网络的结构，用深度学习框架实现
- 训练及改进模型
我自己生成的数据集在GitHub，包含了训练、测试、验证集：
GitHub - wolverinn/HEVC-CU-depths-dataset: A dataset that contains the Coding Unit image files and their corresponding depths for HEVC intra-prediction.
现在我主要做了构建数据集和数据预处理的工作。HEVC的编码器的输入是一帧一帧的图像，输出是图像中的每个CTU的最佳分割深度。因此需要构建的数据集就是图像以及该图像对应的分割信息。
在HEVC部分源码剖析中，通过分析HEVC编码器部分的源代码，我们已经可以对输入的每一帧图像，输出它的CTU分割信息：
向HEVC的编码器送入一个YUV文件，这个YUV文件会被分解成很多帧图像(frame)，每个图像又分成若干个CTU，每个CTU会有一个16x16的分割信息。因此产生数据集的思路就是：使用HEVC的编码器处理YUV文件，就可以得到这个YUV文件的每一帧的CTU的分割信息，再用FFmpeg命令：
ffmpeg -video_size 832x480 -r 50 -pixel_format yuv420p -i BasketballDrill_832x480_50.yuv output-%d.png
将这个YUV文件分解成一帧一帧的图像，就得到了用作神经网络输入的图像。
但是神经网络的训练要求一个很大的数据集，所以显然不能人工一个个地去产生每个YUV文件的CTU分割信息和对应的每一帧图像。需要一个可以批量生成数据的脚本，获取YUV文件，调用HEVC编码器，然后使用FFmpeg从YUV文件中提取出每一帧。并且由于图像和分割信息是单独存放，我们还需要注意命名的问题，确保每张图片和分割信息可以对应。
首先是如何自动化调用HEVC的编码器，先对源代码进行修改，使其能够输出格式化的分割信息到txt文件，然后进行编译。HEVC编码器部分在编译完成后对应的可执行文件是TAppEncoder.exe，它接受-c config_file作为输入参数，我们可以在config_file中指定输入的YUV文件的一些参数如帧率，宽度高度等。然后运行编码器获取输出的txt文件。
自动化调用HEVC编码器并输出格式化的分割信息到txt文件的代码如下：
import os
# this script needs to be in the same directory of the two config files and the encoder.exe: WORKSPACE_PATH
YUV_FILE_PATH = "E:\\HM\\trunk\\workspace\\yuv-resources"
WORKSPACE_PATH = os.getcwd()
CtuInfo_FILENAME = "BasketballdrillCU.txt"
def gen_cfg(yuv_filename):
FrameRate = yuv_filename.split('_')[2].strip(".yuv")
SourceWidth = yuv_filename.split('_')[1].split('x')[0]
SourceHeight = yuv_filename.split('_')[1].split('x')[1]
with open('bitstream.cfg','w') as f:
f.write("InputFile : {}\\{}\n".format(YUV_FILE_PATH,yuv_filename))
f.write("InputBitDepth : 8\n")
f.write("InputChromaFormat : 420\n")
f.write("FrameRate : {}\n".format(FrameRate))
f.write("FrameSkip : 0\n")
f.write("SourceWidth : {}\n".format(SourceWidth))
f.write("SourceHeight : {}\n".format(SourceHeight))
f.write("FramesToBeEncoded : 15000\n")
f.write("Level : 3.1")
encoding_cmd = "TAppEncoder.exe -c encoder_intra_main.cfg -c bitstream.cfg"
for i,yuv_filename in enumerate(os.listdir(YUV_FILE_PATH)):
gen_cfg(yuv_filename)
os.system(encoding_cmd)
os.rename(CtuInfo_FILENAME,"v_{}.txt".format(str(i)))
然后再使用FFmpeg获取YUV文件中的帧，在for循环中加入：
gen_frames_cmd = "ffmpeg -video_size {} -r {} -pixel_format yuv420p -i {}\\{} {}\\img-train\\v_{}_%d_.jpg".format(yuv_filename.split('_')[1],yuv_filename.split('_')[2].strip(".yuv"),YUV_FILE_PATH,yuv_filename,WORKSPACE_PATH,str(i))
os.system(gen_frames_cmd)
这样我们就能得到命名规范的数据集了：
接下来需要考虑数据预处理的问题，需要预处理的数据包括图片数据和分割信息的数据。当我们向神经网络中送入图片时，图片大小最好是2的整数次幂或者一个较大的整数次幂的倍数，这样有利于神经网络的处理。这里我采用的方式是论文 "Fast CU Depth Decision for HEVC Using Neural Networks" 中所使用的每次送入一个CTU大小，也就是64x64 ，所以，使用Pillow库读取一张完整的图片之后，需要定位到此次要提取的CTU，然后将这个CTU裁剪下来，使用Image.crop()函数。
定位CTU及裁剪的图像处理代码如下：
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim
from torch.autograd import Variable
from torchvision import datasets, transforms
import os
import numpy as np
from PIL import Image
import time
import math
BATCH_SIZE=512
EPOCHS=10 # 总共训练批次
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 让torch判断是否使用GPU
IMG_WIDTH = 1920
IMG_HEIGHT = 1080
transform = transforms.Compose([
transforms.ToTensor(), # 将图片转换为Tensor,归一化至[0,1]
# transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]) # 标准化至[-1,1]
])
class ImageSet(data.Dataset):
def __init__(self,root):
# 所有图片的绝对路径
self.img_files = []
self.root = root
for imgs in os.listdir(root):
self.img_files.append(imgs)
self.transforms=transform
def __getitem__(self, index):
ctu_number_per_img = math.ceil(IMG_WIDTH/64)*math.ceil(IMG_HEIGHT/64)
img_index = index//ctu_number_per_img
ctu_number = index%ctu_number_per_img
whole_img = Image.open(os.path.join(self.root,self.img_files[img_index]))
img_row = ctu_number//math.ceil(IMG_WIDTH/64)
img_colonm = ctu_number%math.ceil(IMG_WIDTH/64)
start_pixel_x = (img_colonm-1)*64
start_pixel_y = img_row*64
img = whole_img.crop((start_pixel_x,start_pixel_y,start_pixel_x+64,start_pixel_y+64))
video_number = self.img_files[img_index].split('_')[1]
frame_number = self.img_files[img_index].split('_')[2]
if self.transforms:
data = self.transforms(img)
else:
img = np.asarray(img)
data = torch.from_numpy(img)
label = from_ctufile(video_number,frame_number,str(ctu_number))
return data,label
def __len__(self):
return len(self.img_files)
这里继承了pytorch中的data.Dataset类，是为了训练神经网络加载数据的时候比较方便，加载时直接使用：
train_loader = data.DataLoader(ImageSet("./img-train/"),batch_size=BATCH_SIZE,shuffle=True)
这样，在类ImageSet中的函数__getitem__()返回的“data”就是经过裁剪的图片数据。
接下来还需要处理作为“label”的CTU分割信息文件，每个CTU的分割信息都是16x16的矩阵，由于4x4已经是最小的分割了，因此我们只需要从CTU的分割矩阵中提取出16个分割深度信息就可以了，可以直观地看一下：
在这个矩阵中，每个4x4的小块只需要提取出一个深度信息就够了，最后这个16x16的矩阵可以提取出16个分割深度信息。除此之外，我们还需要确保之前裁剪的CTU和这次提取的CTU是同一个帧内部的同一个CTU，提取label的函数在__getitem__()中的调用是：
label = from_ctufile(video_number,frame_number,str(ctu_number))
使用video_number, frame_number, ctu_number来定位至某个YUV文件下的某一帧的某一个CTU，这个提取label的函数定义为：
def from_ctufile(video_number,frame_number,ctu_number):
ctu_file = "v_{}.txt".format(video_number)
frame_detected = 0
ctu_detected = 0
converting = 0
label_list = []
with open(ctu_file, 'r') as f:
for i, line in enumerate(f):
if ctu_detected == 1 and "ctu" in line:
label = torch.FloatTensor(label_list) # https://pytorch-cn.readthedocs.io/zh/latest/package_references/Tensor/
return label
if frame_detected == 0:
if "frame" in line:
current_frame = line.split(':')[1]
if int(frame_number)-1 == int(current_frame):
frame_detected = 1
else:
continue
else:
continue
elif ctu_detected ==0:
if "ctu" in line:
current_ctu = line.split(':')[1]
if int(ctu_number) == int(current_ctu):
ctu_detected = 1
continue
else:
if (converting) % 4 == 0:
line_depths = line.split(' ')
for index in range(16):
if index % 4 == 0:
label_list.append(line_depths[index])
converting += 1
else:
converting += 1
在含有整个YUV文件的分割信息的txt文件中，先定位frame，再定位CTU，然后进行转换，得到长度为16的label
至此，数据准备工作差不多完成了，还差一步，就是YUV文件的获取，要构建数据集，需要大量的高度和宽度相同的YUV文件，可以先尝试从网上下载现有资源，如果不够的话可以用FFmpeg自己通过其它格式的视频生成
我自己生成的数据集GitHub地址：
