Spaces:
Running
Running
| import torch | |
| import numpy as np | |
| import cv2 | |
| import os | |
| from loss import batch_episym | |
| from tqdm import tqdm | |
| import sys | |
| ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| sys.path.insert(0, ROOT_DIR) | |
| from utils import evaluation_utils, train_utils | |
| def valid(valid_loader, model, match_loss, config, model_config): | |
| model.eval() | |
| loader_iter = iter(valid_loader) | |
| num_pair = 0 | |
| total_loss, total_acc_corr, total_acc_incorr = 0, 0, 0 | |
| total_precision, total_recall = torch.zeros( | |
| model_config.layer_num, device="cuda" | |
| ), torch.zeros(model_config.layer_num, device="cuda") | |
| total_acc_mid = torch.zeros(len(model_config.seedlayer) - 1, device="cuda") | |
| with torch.no_grad(): | |
| if config.local_rank == 0: | |
| loader_iter = tqdm(loader_iter) | |
| print("validating...") | |
| for test_data in loader_iter: | |
| num_pair += 1 | |
| test_data = train_utils.tocuda(test_data) | |
| res = model(test_data) | |
| loss_res = match_loss.run(test_data, res) | |
| total_acc_corr += loss_res["acc_corr"] | |
| total_acc_incorr += loss_res["acc_incorr"] | |
| total_loss += loss_res["total_loss"] | |
| if config.model_name == "SGM": | |
| total_acc_mid += loss_res["mid_acc_corr"] | |
| total_precision, total_recall = ( | |
| total_precision + loss_res["pre_seed_conf"], | |
| total_recall + loss_res["recall_seed_conf"], | |
| ) | |
| total_acc_corr /= num_pair | |
| total_acc_incorr /= num_pair | |
| total_precision /= num_pair | |
| total_recall /= num_pair | |
| total_acc_mid /= num_pair | |
| # apply tensor reduction | |
| ( | |
| total_loss, | |
| total_acc_corr, | |
| total_acc_incorr, | |
| total_precision, | |
| total_recall, | |
| total_acc_mid, | |
| ) = ( | |
| train_utils.reduce_tensor(total_loss, "sum"), | |
| train_utils.reduce_tensor(total_acc_corr, "mean"), | |
| train_utils.reduce_tensor(total_acc_incorr, "mean"), | |
| train_utils.reduce_tensor(total_precision, "mean"), | |
| train_utils.reduce_tensor(total_recall, "mean"), | |
| train_utils.reduce_tensor(total_acc_mid, "mean"), | |
| ) | |
| model.train() | |
| return ( | |
| total_loss, | |
| total_acc_corr, | |
| total_acc_incorr, | |
| total_precision, | |
| total_recall, | |
| total_acc_mid, | |
| ) | |
| def dump_train_vis(res, data, step, config): | |
| # batch matching | |
| p = res["p"][:, :-1, :-1] | |
| score, index1 = torch.max(p, dim=-1) | |
| _, index2 = torch.max(p, dim=-2) | |
| mask_th = score > 0.2 | |
| mask_mc = index2.gather(index=index1, dim=1) == torch.arange(len(p[0])).cuda()[None] | |
| mask_p = mask_th & mask_mc # B*N | |
| corr1, corr2 = data["x1"], data["x2"].gather( | |
| index=index1[:, :, None].expand(-1, -1, 2), dim=1 | |
| ) | |
| corr1_kpt, corr2_kpt = data["kpt1"], data["kpt2"].gather( | |
| index=index1[:, :, None].expand(-1, -1, 2), dim=1 | |
| ) | |
| epi_dis = batch_episym(corr1, corr2, data["e_gt"]) | |
| mask_inlier = epi_dis < config.inlier_th # B*N | |
| # dump vis | |
| for cur_mask_p, cur_mask_inlier, cur_corr1, cur_corr2, img_path1, img_path2 in zip( | |
| mask_p, mask_inlier, corr1_kpt, corr2_kpt, data["img_path1"], data["img_path2"] | |
| ): | |
| img1, img2 = cv2.imread(img_path1), cv2.imread(img_path2) | |
| dis_play = evaluation_utils.draw_match( | |
| img1, | |
| img2, | |
| cur_corr1[cur_mask_p].cpu().numpy(), | |
| cur_corr2[cur_mask_p].cpu().numpy(), | |
| inlier=cur_mask_inlier, | |
| ) | |
| base_name_seq = os.path.join( | |
| img_path1.split("/")[-1] | |
| + "_" | |
| + img_path2.split("/")[-1] | |
| + "_" | |
| + img_path1.split("/")[-2] | |
| ) | |
| save_path = os.path.join( | |
| config.train_vis_folder, | |
| "train_vis", | |
| config.log_base, | |
| str(step), | |
| base_name_seq + ".png", | |
| ) | |
| cv2.imwrite(save_path, dis_play) | |