以下内容参考: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