def flatten_64to1(semi, cell_size=8):
"""
input:
semi: tensor[batch, cell_size*cell_size, Hc, Wc]
(Hc = H/8)
outpus:
heatmap: tensor[batch, 1, H, W]
"""
from utils.d2s import DepthToSpace
depth2space = DepthToSpace(cell_size)
heatmap = depth2space(semi)
return heatmap
def flattenDetection(semi, tensor=False):
'''
Flatten detection output
:param semi:
output from detector head
tensor [65, Hc, Wc]
:or
tensor (batch_size, 65, Hc, Wc)
:return:
3D heatmap
np (1, H, C)
:or
tensor (batch_size, 65, Hc, Wc)
'''
batch = False
if len(semi.shape) == 4:
batch = True
batch_size = semi.shape[0]
# if tensor:
# semi.exp_()
# d = semi.sum(dim=1) + 0.00001
# d = d.view(d.shape[0], 1, d.shape[1], d.shape[2])
# semi = semi / d # how to /(64,15,20)
# nodust = semi[:, :-1, :, :]
# heatmap = flatten64to1(nodust, tensor=tensor)
# else:
# Convert pytorch -> numpy.
# --- Process points.
# dense = nn.functional.softmax(semi, dim=0) # [65, Hc, Wc]
if batch:
dense = nn.functional.softmax(semi, dim=1) # [batch, 65, Hc, Wc]
# Remove dustbin.
nodust = dense[:, :-1, :, :]
else:
dense = nn.functional.softmax(semi, dim=0) # [65, Hc, Wc]
nodust = dense[:-1, :, :].unsqueeze(0)
# Reshape to get full resolution heatmap.
# heatmap = flatten64to1(nodust, tensor=True) # [1, H, W]
depth2space = DepthToSpace(8)
heatmap = depth2space(nodust)
heatmap = heatmap.squeeze(0) if not batch else heatmap
return heatmap