def update_overlap(images_dict, labels_warp_2D, heatmap_nms_batch,
img_warp, name):
# image overlap
from utils.draw import img_overlap
# result_overlap = img_overlap(img_r, img_g, img_gray)
# overlap label, nms, img
nms_overlap = [
img_overlap(
toNumpy(labels_warp_2D[i]),
heatmap_nms_batch[i],
toNumpy(img_warp[i]),
) for i in range(heatmap_nms_batch.shape[0])
]
nms_overlap = np.stack(nms_overlap, axis=0)
images_dict.update({name + "_nms_overlap": nms_overlap})
def run(self, images):
"""
input:
images: tensor[batch(1), 1, H, W]
"""
from Train_model_heatmap import Train_model_heatmap
from utils.var_dim import toNumpy
train_agent = Train_model_heatmap
with torch.no_grad():
outs = self.net(images)
semi = outs['semi']
self.outs = outs
channel = semi.shape[1]
if channel == 64:
heatmap = train_agent.flatten_64to1(semi, cell_size=self.cell_size)
elif channel == 65:
heatmap = flattenDetection(semi, tensor=True)
heatmap_np = toNumpy(heatmap)
self.heatmap = heatmap_np
return heatmap
pass
def heatmap_to_nms(self, heatmap, tensor=False, boxnms=False):
"""
return:
heatmap_nms_batch: np [batch, 1, H, W]
"""
to_floatTensor = lambda x: torch.from_numpy(x).type(torch.FloatTensor)
from utils.var_dim import toNumpy
heatmap_np = toNumpy(heatmap)
## heatmap_nms
if boxnms:
from utils.utils import box_nms
heatmap_nms_batch = [box_nms(h.detach().squeeze(), self.nms_dist, min_prob=self.conf_thresh) \
for h in heatmap] # [batch, H, W]
heatmap_nms_batch = torch.stack(heatmap_nms_batch,
dim=0).unsqueeze(1)
# print('heatmap_nms_batch: ', heatmap_nms_batch.shape)
else:
heatmap_nms_batch = [self.heatmap_nms(h, self.nms_dist, self.conf_thresh) \
for h in heatmap_np] # [batch, H, W]
heatmap_nms_batch = np.stack(heatmap_nms_batch, axis=0)
heatmap_nms_batch = heatmap_nms_batch[:, np.newaxis, ...]
if tensor:
heatmap_nms_batch = to_floatTensor(heatmap_nms_batch)
heatmap_nms_batch = heatmap_nms_batch.to(self.device)
self.heatmap = heatmap
self.heatmap_nms_batch = heatmap_nms_batch
return heatmap_nms_batch
pass
def heatmap_to_nms(self, images_dict, heatmap, name):
"""
return:
heatmap_nms_batch: np [batch, H, W]
"""
from utils.var_dim import toNumpy
heatmap_np = toNumpy(heatmap)
## heatmap_nms
heatmap_nms_batch = [self.heatmap_nms(h) for h in heatmap_np] # [batch, H, W]
heatmap_nms_batch = np.stack(heatmap_nms_batch, axis=0)
# images_dict.update({name + '_nms_batch': heatmap_nms_batch})
images_dict.update({name + "_nms_batch": heatmap_nms_batch[:, np.newaxis, ...]})
return heatmap_nms_batch
def train_val_sample(self, sample, n_iter=0, train=False):
"""
# key function
:param sample:
:param n_iter:
:param train:
:return:
"""
to_floatTensor = lambda x: torch.tensor(x).type(torch.FloatTensor)
task = "train" if train else "val"
tb_interval = self.config["tensorboard_interval"]
if_warp = self.config['data']['warped_pair']['enable']
self.scalar_dict, self.images_dict, self.hist_dict = {}, {}, {}
## get the inputs
# logging.info('get input img and label')
img, labels_2D, mask_2D = (
sample["image"],
sample["labels_2D"],
sample["valid_mask"],
)
# img, labels = img.to(self.device), labels_2D.to(self.device)
# variables
batch_size, H, W = img.shape[0], img.shape[2], img.shape[3]
self.batch_size = batch_size
det_loss_type = self.config["model"]["detector_loss"]["loss_type"]
# print("batch_size: ", batch_size)
Hc = H // self.cell_size
Wc = W // self.cell_size
# warped images
# img_warp, labels_warp_2D, mask_warp_2D = sample['warped_img'].to(self.device), \
# sample['warped_labels'].to(self.device), \
# sample['warped_valid_mask'].to(self.device)
if if_warp:
img_warp, labels_warp_2D, mask_warp_2D = (
sample["warped_img"],
sample["warped_labels"],
sample["warped_valid_mask"],
)
# homographies
# mat_H, mat_H_inv = \
# sample['homographies'].to(self.device), sample['inv_homographies'].to(self.device)
if if_warp:
mat_H, mat_H_inv = sample["homographies"], sample[
"inv_homographies"]
# zero the parameter gradients
self.optimizer.zero_grad()
# forward + backward + optimize
if train:
# print("img: ", img.shape, ", img_warp: ", img_warp.shape)
outs = self.net(img.to(self.device))
semi, coarse_desc = outs["semi"], outs["desc"]
if if_warp:
outs_warp = self.net(img_warp.to(self.device))
semi_warp, coarse_desc_warp = outs_warp["semi"], outs_warp[
"desc"]
else:
with torch.no_grad():
outs = self.net(img.to(self.device))
semi, coarse_desc = outs["semi"], outs["desc"]
if if_warp:
outs_warp = self.net(img_warp.to(self.device))
semi_warp, coarse_desc_warp = outs_warp["semi"], outs_warp[
"desc"]
pass
# detector loss
from utils.utils import labels2Dto3D
if self.gaussian:
labels_2D = sample["labels_2D_gaussian"]
if if_warp:
warped_labels = sample["warped_labels_gaussian"]
else:
labels_2D = sample["labels_2D"]
if if_warp:
warped_labels = sample["warped_labels"]
add_dustbin = False
if det_loss_type == "l2":
add_dustbin = False
elif det_loss_type == "softmax":
add_dustbin = True
labels_3D = labels2Dto3D(labels_2D.to(self.device),
cell_size=self.cell_size,
add_dustbin=add_dustbin).float()
mask_3D_flattened = self.getMasks(mask_2D,
self.cell_size,
device=self.device)
loss_det = self.detector_loss(
input=outs["semi"],
target=labels_3D.to(self.device),
mask=mask_3D_flattened,
loss_type=det_loss_type,
)
# warp
if if_warp:
labels_3D = labels2Dto3D(
warped_labels.to(self.device),
cell_size=self.cell_size,
add_dustbin=add_dustbin,
).float()
mask_3D_flattened = self.getMasks(mask_warp_2D,
self.cell_size,
device=self.device)
loss_det_warp = self.detector_loss(
input=outs_warp["semi"],
target=labels_3D.to(self.device),
mask=mask_3D_flattened,
loss_type=det_loss_type,
)
else:
loss_det_warp = torch.tensor([0]).to(self.device)
## get labels, masks, loss for detection
# labels3D_in_loss = self.getLabels(labels_2D, self.cell_size, device=self.device)
# mask_3D_flattened = self.getMasks(mask_2D, self.cell_size, device=self.device)
# loss_det = self.get_loss(semi, labels3D_in_loss, mask_3D_flattened, device=self.device)
## warping
# labels3D_in_loss = self.getLabels(labels_warp_2D, self.cell_size, device=self.device)
# mask_3D_flattened = self.getMasks(mask_warp_2D, self.cell_size, device=self.device)
# loss_det_warp = self.get_loss(semi_warp, labels3D_in_loss, mask_3D_flattened, device=self.device)
mask_desc = mask_3D_flattened.unsqueeze(1)
lambda_loss = self.config["model"]["lambda_loss"]
# print("mask_desc: ", mask_desc.shape)
# print("mask_warp_2D: ", mask_warp_2D.shape)
# descriptor loss
if lambda_loss > 0:
assert if_warp == True, "need a pair of images"
loss_desc, mask, positive_dist, negative_dist = self.descriptor_loss(
coarse_desc,
coarse_desc_warp,
mat_H,
mask_valid=mask_desc,
device=self.device,
**self.desc_params)
else:
ze = torch.tensor([0]).to(self.device)
loss_desc, positive_dist, negative_dist = ze, ze, ze
loss = loss_det + loss_det_warp
if lambda_loss > 0:
loss += lambda_loss * loss_desc
##### try to minimize the error ######
add_res_loss = False
if add_res_loss and n_iter % 10 == 0:
print("add_res_loss!!!")
heatmap_org = self.get_heatmap(semi, det_loss_type) # tensor []
heatmap_org_nms_batch = self.heatmap_to_nms(self.images_dict,
heatmap_org,
name="heatmap_org")
if if_warp:
heatmap_warp = self.get_heatmap(semi_warp, det_loss_type)
heatmap_warp_nms_batch = self.heatmap_to_nms(
self.images_dict, heatmap_warp, name="heatmap_warp")
# original: pred
## check the loss on given labels!
outs_res = self.get_residual_loss(
sample["labels_2D"] *
to_floatTensor(heatmap_org_nms_batch).unsqueeze(1),
heatmap_org,
sample["labels_res"],
name="original_pred",
)
loss_res_ori = (outs_res["loss"]**2).mean()
# warped: pred
if if_warp:
outs_res_warp = self.get_residual_loss(
sample["warped_labels"] *
to_floatTensor(heatmap_warp_nms_batch).unsqueeze(1),
heatmap_warp,
sample["warped_res"],
name="warped_pred",
)
loss_res_warp = (outs_res_warp["loss"]**2).mean()
else:
loss_res_warp = torch.tensor([0]).to(self.device)
loss_res = loss_res_ori + loss_res_warp
# print("loss_res requires_grad: ", loss_res.requires_grad)
loss += loss_res
self.scalar_dict.update({
"loss_res_ori": loss_res_ori,
"loss_res_warp": loss_res_warp
})
#######################################
self.loss = loss
self.scalar_dict.update({
"loss": loss,
"loss_det": loss_det,
"loss_det_warp": loss_det_warp,
"positive_dist": positive_dist,
"negative_dist": negative_dist,
})
self.input_to_imgDict(sample, self.images_dict)
if train:
loss.backward()
self.optimizer.step()
if n_iter % tb_interval == 0 or task == "val":
logging.info("current iteration: %d, tensorboard_interval: %d",
n_iter, tb_interval)
# add clean map to tensorboard
## semi_warp: flatten, to_numpy
heatmap_org = self.get_heatmap(semi, det_loss_type) # tensor []
heatmap_org_nms_batch = self.heatmap_to_nms(self.images_dict,
heatmap_org,
name="heatmap_org")
if if_warp:
heatmap_warp = self.get_heatmap(semi_warp, det_loss_type)
heatmap_warp_nms_batch = self.heatmap_to_nms(
self.images_dict, heatmap_warp, name="heatmap_warp")
def update_overlap(images_dict, labels_warp_2D, heatmap_nms_batch,
img_warp, name):
# image overlap
from utils.draw import img_overlap
# result_overlap = img_overlap(img_r, img_g, img_gray)
# overlap label, nms, img
nms_overlap = [
img_overlap(
toNumpy(labels_warp_2D[i]),
heatmap_nms_batch[i],
toNumpy(img_warp[i]),
) for i in range(heatmap_nms_batch.shape[0])
]
nms_overlap = np.stack(nms_overlap, axis=0)
images_dict.update({name + "_nms_overlap": nms_overlap})
from utils.var_dim import toNumpy
update_overlap(
self.images_dict,
labels_2D,
heatmap_org_nms_batch[np.newaxis, ...],
img,
"original",
)
update_overlap(
self.images_dict,
labels_2D,
toNumpy(heatmap_org),
img,
"original_heatmap",
)
if if_warp:
update_overlap(
self.images_dict,
labels_warp_2D,
heatmap_warp_nms_batch[np.newaxis, ...],
img_warp,
"warped",
)
update_overlap(
self.images_dict,
labels_warp_2D,
toNumpy(heatmap_warp),
img_warp,
"warped_heatmap",
)
# residuals
from utils.losses import do_log
if self.gaussian:
# original: gt
self.get_residual_loss(
sample["labels_2D"],
sample["labels_2D_gaussian"],
sample["labels_res"],
name="original_gt",
)
if if_warp:
# warped: gt
self.get_residual_loss(
sample["warped_labels"],
sample["warped_labels_gaussian"],
sample["warped_res"],
name="warped_gt",
)
# from utils.losses import do_log
# patches_log = do_log(patches)
# original: pred
## check the loss on given labels!
# self.get_residual_loss(
# sample["labels_2D"]
# * to_floatTensor(heatmap_org_nms_batch).unsqueeze(1),
# heatmap_org,
# sample["labels_res"],
# name="original_pred",
# )
# print("heatmap_org_nms_batch: ", heatmap_org_nms_batch.shape)
# get_residual_loss(to_floatTensor(heatmap_org_nms_batch).unsqueeze(1), heatmap_org,
# sample['labels_res'], name='original_pred')
# warped: pred
# self.get_residual_loss(
# sample["warped_labels"]
# * to_floatTensor(heatmap_warp_nms_batch).unsqueeze(1),
# heatmap_warp,
# sample["warped_res"],
# name="warped_pred",
# )
# get_residual_loss(to_floatTensor(heatmap_warp_nms_batch).unsqueeze(1), heatmap_warp,
# sample['warped_res'], name='warped_pred')
# precision, recall
# pr_mean = self.batch_precision_recall(
# to_floatTensor(heatmap_warp_nms_batch[:, np.newaxis, ...]),
# sample["warped_labels"],
# )
pr_mean = self.batch_precision_recall(
to_floatTensor(heatmap_org_nms_batch[:, np.newaxis, ...]),
sample["labels_2D"],
)
print("pr_mean")
self.scalar_dict.update(pr_mean)
self.printLosses(self.scalar_dict, task)
self.tb_images_dict(task, self.images_dict, max_img=2)
self.tb_hist_dict(task, self.hist_dict)
self.tb_scalar_dict(self.scalar_dict, task)
return loss.item()
