HEVC自帶的編碼器在決定CTU的最佳分割深度的時候會花費(fèi)很多時間盏求，我們的目標(biāo)是訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)來根據(jù)輸入的幀預(yù)測這一幀圖像上面的所有CTU的最佳分割深度，這樣會比HEVC自帶的編碼器快很多亿眠。

我們需要做的主要有：

• 構(gòu)建用于訓(xùn)練神經(jīng)網(wǎng)絡(luò)的數(shù)據(jù)集
• 數(shù)據(jù)預(yù)處理
• 構(gòu)建卷積神經(jīng)網(wǎng)絡(luò)的結(jié)構(gòu)碎罚，用深度學(xué)習(xí)框架實(shí)現(xiàn)
• 訓(xùn)練及改進(jìn)模型

我自己生成的數(shù)據(jù)集在GitHub，包含了訓(xùn)練纳像、測試荆烈、驗(yàn)證集：

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

現(xiàn)在我主要做了構(gòu)建數(shù)據(jù)集和數(shù)據(jù)預(yù)處理的工作。HEVC的編碼器的輸入是一幀一幀的圖像竟趾，輸出是圖像中的每個CTU的最佳分割深度耙考。因此需要構(gòu)建的數(shù)據(jù)集就是圖像以及該圖像對應(yīng)的分割信息。

在HEVC部分源碼剖析中潭兽，通過分析HEVC編碼器部分的源代碼，我們已經(jīng)可以對輸入的每一幀圖像斗遏，輸出它的CTU分割信息：

PartitionInfo

向HEVC的編碼器送入一個YUV文件山卦，這個YUV文件會被分解成很多幀圖像(frame)，每個圖像又分成若干個CTU，每個CTU會有一個16x16的分割信息账蓉。因此產(chǎn)生數(shù)據(jù)集的思路就是：使用HEVC的編碼器處理YUV文件枚碗，就可以得到這個YUV文件的每一幀的CTU的分割信息，再用FFmpeg命令：

ffmpeg -video_size 832x480 -r 50 -pixel_format yuv420p -i BasketballDrill_832x480_50.yuv output-%d.png

將這個YUV文件分解成一幀一幀的圖像铸本，就得到了用作神經(jīng)網(wǎng)絡(luò)輸入的圖像肮雨。

但是神經(jīng)網(wǎng)絡(luò)的訓(xùn)練要求一個很大的數(shù)據(jù)集，所以顯然不能人工一個個地去產(chǎn)生每個YUV文件的CTU分割信息和對應(yīng)的每一幀圖像箱玷。需要一個可以批量生成數(shù)據(jù)的腳本怨规，獲取YUV文件，調(diào)用HEVC編碼器锡足，然后使用FFmpeg從YUV文件中提取出每一幀波丰。并且由于圖像和分割信息是單獨(dú)存放，我們還需要注意命名的問題舶得，確保每張圖片和分割信息可以對應(yīng)掰烟。

首先是如何自動化調(diào)用HEVC的編碼器，先對源代碼進(jìn)行修改沐批，使其能夠輸出格式化的分割信息到txt文件纫骑，然后進(jìn)行編譯。HEVC編碼器部分在編譯完成后對應(yīng)的可執(zhí)行文件是TAppEncoder.exe九孩，它接受-c config_file作為輸入?yún)?shù)先馆，我們可以在config_file中指定輸入的YUV文件的一些參數(shù)如幀率，寬度高度等捻撑。然后運(yùn)行編碼器獲取輸出的txt文件磨隘。

自動化調(diào)用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循環(huán)中加入：

    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)

這樣我們就能得到命名規(guī)范的數(shù)據(jù)集了：

dataset

接下來需要考慮數(shù)據(jù)預(yù)處理的問題顾患，需要預(yù)處理的數(shù)據(jù)包括圖片數(shù)據(jù)和分割信息的數(shù)據(jù)番捂。當(dāng)我們向神經(jīng)網(wǎng)絡(luò)中送入圖片時，圖片大小最好是2的整數(shù)次冪或者一個較大的整數(shù)次冪的倍數(shù)江解，這樣有利于神經(jīng)網(wǎng)絡(luò)的處理设预。這里我采用的方式是論文 "Fast CU Depth Decision for HEVC Using Neural Networks" 中所使用的每次送入一個CTU大小，也就是64x64 犁河，所以鳖枕，使用Pillow庫讀取一張完整的圖片之后，需要定位到此次要提取的CTU桨螺，然后將這個CTU裁剪下來宾符，使用Image.crop()函數(shù)。

定位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 # 總共訓(xùn)練批次
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 讓torch判斷是否使用GPU
IMG_WIDTH = 1920
IMG_HEIGHT = 1080
transform = transforms.Compose([
    transforms.ToTensor(),  # 將圖片轉(zhuǎn)換為Tensor,歸一化至[0,1]
    # transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])  # 標(biāo)準(zhǔn)化至[-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類灭翔，是為了訓(xùn)練神經(jīng)網(wǎng)絡(luò)加載數(shù)據(jù)的時候比較方便魏烫，加載時直接使用：

train_loader = data.DataLoader(ImageSet("./img-train/"),batch_size=BATCH_SIZE,shuffle=True)

這樣，在類ImageSet中的函數(shù)__getitem__()返回的“data”就是經(jīng)過裁剪的圖片數(shù)據(jù)。

接下來還需要處理作為“l(fā)abel”的CTU分割信息文件哄褒，每個CTU的分割信息都是16x16的矩陣稀蟋，由于4x4已經(jīng)是最小的分割了，因此我們只需要從CTU的分割矩陣中提取出16個分割深度信息就可以了呐赡，可以直觀地看一下：

one CTU

在這個矩陣中退客，每個4x4的小塊只需要提取出一個深度信息就夠了，最后這個16x16的矩陣可以提取出16個分割深度信息链嘀。除此之外萌狂，我們還需要確保之前裁剪的CTU和這次提取的CTU是同一個幀內(nèi)部的同一個CTU，提取label的函數(shù)在__getitem__()中的調(diào)用是：

label = from_ctufile(video_number,frame_number,str(ctu_number))

使用video_number, frame_number, ctu_number來定位至某個YUV文件下的某一幀的某一個CTU管闷，這個提取label的函數(shù)定義為：

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包个，然后進(jìn)行轉(zhuǎn)換刷允，得到長度為16的label

至此，數(shù)據(jù)準(zhǔn)備工作差不多完成了碧囊，還差一步树灶，就是YUV文件的獲取，要構(gòu)建數(shù)據(jù)集糯而，需要大量的高度和寬度相同的YUV文件天通，可以先嘗試從網(wǎng)上下載現(xiàn)有資源，如果不夠的話可以用FFmpeg自己通過其它格式的視頻生成

我自己生成的數(shù)據(jù)集GitHub地址：

HEVC-CU-depths-dataset