HM16.9 提取CU depth，cu partition可视化

以下内容参考：https://github.com/pengrenlai/HEVC-CUpartition-visualization

提取CU depth，cu partition可视化(hm16.9)

1，在TEncGOP.cpp 中的m_pcLoopFilter->loopFilterPic( pcPic );下面添加代码提取cu depth。遍历当前编码帧所有的CTU，再遍历CTU中每个4x4的depth block，按照Raster scan order 输出到文件中。

    m_pcLoopFilter->loopFilterPic( pcPic );
 start add code here ///
    ofstream myfile_CU;
	string CU_file_name = "cu_depth.txt";
	myfile_CU.open(CU_file_name);

for (UInt ctuRsAddr = 0; ctuRsAddr < pcPic->getNumberOfCtusInFrame(); ctuRsAddr++)
{
    TComDataCU* pCtu = pcPic->getCtu(ctuRsAddr);
    for (int i = 0; i < pCtu->getTotalNumPart(); i++)
    {
        myfile_CU << to_string(pCtu->getDepth(g_auiRasterToZscan[i])) + " ";
    }
}
myfile_CU.close();
/end

2，用python 从文件中读取cu depth。把每个cu-depth block 映射到 4x4 像素上,这样就可以得到像素级的depth_map. 然后需要做的就是将depth_map 转换成partition map。如何把这个partition map画出来呢？解决这个问题的思想是将像素的位置与depth相结合来得到partition map。比如当前像素的位置是[i,j],深度是d[i,j] 。如果d[i.j]是0，而且i，或j是64的倍数，那么说明这个像素一定是cu的边界。因为四叉树划分深度0对应的是64x64的像素块。类似可以推出深度是1，2，3的像素块的边界。主要实现是下面这两个函数

def isVerBoundary(i, j, CUorTU_size, CUorTU_depth):
        left = (j+1) % CUorTU_size == 0 and (j+1) < widthInCTU*64 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i][j+1]) >= CUorTU_depth)
        right = j % CUorTU_size == 0 and j-1 >= 0 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i][j-1]) >= CUorTU_depth)
        return left or right

    def isHorBoundary(i, j, CUorTU_size, CUorTU_depth):
        up = (i+1) % CUorTU_size == 0 and (i+1) < heightInCTU*64 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i+1][j]) >= CUorTU_depth)
        down = i % CUorTU_size == 0 and i-1 >= 0 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i-1][j]) >= CUorTU_depth)
        return up or down

3，全部python代码

def save_png(Y,name):   
    Y = Y.astype(np.uint8)
    # Y= Y[0,0,:,:]
    Image.fromarray(Y, mode='L').save(name) 


from PIL import Image
import numpy as np
aspect_ratio={
    'ClassB_shape':(1920, 1080), 'ClassB_shape_pad':(1920, 1088),
    'ClassC_shape':(832, 480), 'ClassC_shape_pad':(832, 512),
    'ClassD_shape':(416, 240), 'ClassD_shape_pad':(448, 256),
    'ClassE_shape':(1280, 720), 'ClassE_shape_pad':(1280, 768)
}

w = None
h = None

w_pad = None
h_pad = None

widthInCTU = None
heightInCTU = None

#get CU and TU depth map list. we can leverage it to generate CU and TU partiition map latter.
def get_depth_map_list(depth_stream):
    #CTU row
    row = 0
    #every line in CTU
    line = 0
    count = 0

    curCTU = 0

    depth_map_list = []
    for i in range(heightInCTU*64):
        depth_map_list.append([])

    for d in depth_stream:
        for l in range(4):
            for _ in range(4):
                depth_map_list[row*64+line*4+l].append(d)
        count += 16
        #a CTU is done. update some values
        if count == 4096:
            line = 0
            count = 0
            curCTU += 1
            if curCTU % widthInCTU == 0:
                row += 1
            continue
        #a line in a CTU is done. go to next line
        if count%256 == 0:
            line += 1
    return depth_map_list

# '1' : pixels of edges, '0' : otherwise
def get_partition_map(depth_map_list):
    parition_map_list =[]
    for i in range(heightInCTU*64):
        parition_map_list.append([])

    def isVerBoundary(i, j, CUorTU_size, CUorTU_depth):
        left = (j+1) % CUorTU_size == 0 and (j+1) < widthInCTU*64 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i][j+1]) >= CUorTU_depth)
        right = j % CUorTU_size == 0 and j-1 >= 0 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i][j-1]) >= CUorTU_depth)
        return left or right

    def isHorBoundary(i, j, CUorTU_size, CUorTU_depth):
        up = (i+1) % CUorTU_size == 0 and (i+1) < heightInCTU*64 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i+1][j]) >= CUorTU_depth)
        down = i % CUorTU_size == 0 and i-1 >= 0 and (int(depth_map_list[i][j]) >= CUorTU_depth and int(depth_map_list[i-1][j]) >= CUorTU_depth)
        return up or down

    for i in range(heightInCTU*64):
        for j in range(widthInCTU*64):
            if isVerBoundary(i, j, 64, 0) or isVerBoundary(i, j, 32, 1) or isVerBoundary(i, j, 16, 2) or isVerBoundary(i, j, 8, 3) or isVerBoundary(i, j, 4, 4):
                parition_map_list[i].append(1) #pixels of edges
            elif isHorBoundary(i, j, 64, 0) or isHorBoundary(i, j, 32, 1) or isHorBoundary(i, j, 16, 2) or isHorBoundary(i, j, 8, 3) or isHorBoundary(i, j, 4, 4):
                parition_map_list[i].append(1) #pixels of edges
            else:
                parition_map_list[i].append(0)

    parition_map = np.array(parition_map_list)
    return parition_map

def get_partition_map_from_frame_seq(file_name):        

    frame_file = open(file_name, 'r')
    string = frame_file.read()
    string = string.split()
    depth_map_list = get_depth_map_list(string)
    partition_map = get_partition_map(depth_map_list)

    # img = Image.fromarray(partition_map.astype('uint8')*255).convert('YCbCr') #*255 for visualization => edge:white otherwise:black

    return partition_map

#test
if __name__=='__main__':
    
    w = aspect_ratio['ClassB_shape'][0]
    h = aspect_ratio['ClassB_shape'][1]
    w_pad = aspect_ratio['ClassB_shape_pad'][0]
    h_pad = aspect_ratio['ClassB_shape_pad'][1]
    widthInCTU = w_pad//64
    heightInCTU = h_pad//64
    cu_map = get_partition_map_from_frame_seq('/home/yuejian/HM_CU/workspace/kimol/cu_kimol.txt')
    np.save('kimol_cumap',cu_map[0:1080,:])

    cu_map=cu_map.astype('uint8')*255
    print(cu_map.shape)
    save_png(cu_map,'kimol_cumap.png')

4,可视化

4.1 原图

4.2 通过码流分析工具得到的分割图

4.3 我们最后得到的cu_patition_map

参考：https://github.com/pengrenlai/HEVC-CUpartition-visualization