def pred_soft_argmax(labels_2D, heatmap, labels_res, patch_size=5, device="cuda"):
"""
return:
dict {'loss': mean of difference btw pred and res}
"""
from utils.losses import norm_patches
outs = {}
# extract patches
from utils.losses import extract_patches
from utils.losses import soft_argmax_2d
label_idx = labels_2D[...].nonzero().long()
# patch_size = self.config['params']['patch_size']
patches = extract_patches(
label_idx.to(device), heatmap.to(device), patch_size=patch_size
)
# norm patches
patches = norm_patches(patches)
# predict offsets
from utils.losses import do_log
patches_log = do_log(patches)
# soft_argmax
dxdy = soft_argmax_2d(
patches_log, normalized_coordinates=False
) # tensor [B, N, patch, patch]
dxdy = dxdy.squeeze(1) # tensor [N, 2]
dxdy = dxdy - patch_size // 2
# extract residual
def ext_from_points(labels_res, points):
"""
input:
labels_res: tensor [batch, channel, H, W]
points: tensor [N, 4(pos0(batch), pos1(0), pos2(H), pos3(W) )]
return:
tensor [N, channel]
"""
labels_res = labels_res.transpose(1, 2).transpose(2, 3).unsqueeze(1)
points_res = labels_res[
points[:, 0], points[:, 1], points[:, 2], points[:, 3], :
] # tensor [N, 2]
return points_res
points_res = ext_from_points(labels_res, label_idx)
# loss
outs["pred"] = dxdy
outs["points_res"] = points_res
# ls = lambda x, y: dxdy.cpu() - points_res.cpu()
# outs['loss'] = dxdy.cpu() - points_res.cpu()
outs["loss"] = dxdy.to(device) - points_res.to(device)
outs["patches"] = patches
return outs
def soft_argmax_points(self, pts, patch_size=5):
"""
input:
pts: tensor [N x 2]
"""
from utils.utils import toNumpy
from utils.losses import extract_patch_from_points
from utils.losses import soft_argmax_2d
from utils.losses import norm_patches
##### check not take care of batch #####
# print("not take care of batch! only take first element!")
pts = pts[0].transpose().copy()
patches = extract_patch_from_points(self.heatmap,
pts,
patch_size=patch_size)
import torch
patches = np.stack(patches)
patches_torch = torch.tensor(patches, dtype=torch.float32).unsqueeze(0)
# norm patches
patches_torch = norm_patches(patches_torch)
from utils.losses import do_log
patches_torch = do_log(patches_torch)
# patches_torch = do_log(patches_torch)
# print("one tims of log!")
# print("patches: ", patches_torch.shape)
# print("pts: ", pts.shape)
dxdy = soft_argmax_2d(patches_torch, normalized_coordinates=False)
# print("dxdy: ", dxdy.shape)
points = pts
points[:, :2] = points[:, :2] + dxdy.numpy().squeeze(
) - patch_size // 2
self.patches = patches_torch.numpy().squeeze()
self.pts_subpixel = [points.transpose().copy()]
return self.pts_subpixel.copy()