vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/keypt2subpx/submodules/glue_factory/tests/test_eval_utils.py

@@ -0,0 +1,88 @@
+import unittest
+
+import torch
+
+from gluefactory.eval.utils import eval_matches_homography
+from gluefactory.geometry.homography import warp_points_torch
+
+
+class TestEvalUtils(unittest.TestCase):
+    @staticmethod
+    def default_pts():
+        return torch.tensor(
+            [
+                [10.0, 10.0],
+                [10.0, 20.0],
+                [20.0, 20.0],
+                [20.0, 10.0],
+            ]
+        )
+
+    @staticmethod
+    def default_pred(kps0, kps1):
+        return {
+            "keypoints0": kps0,
+            "keypoints1": kps1,
+            "matches0": torch.arange(len(kps0)),
+            "matching_scores0": torch.ones(len(kps1)),
+        }
+
+    def test_eval_matches_homography_trivial(self):
+        data = {"H_0to1": torch.eye(3)}
+        kps = self.default_pts()
+        pred = self.default_pred(kps, kps)
+
+        results = eval_matches_homography(data, pred)
+
+        self.assertEqual(results["prec@1px"], 1)
+        self.assertEqual(results["prec@3px"], 1)
+        self.assertEqual(results["num_matches"], 4)
+        self.assertEqual(results["num_keypoints"], 4)
+
+    def test_eval_matches_homography_real(self):
+        data = {"H_0to1": torch.tensor([[1.5, 0.2, 21], [-0.3, 1.6, 33], [0, 0, 1.0]])}
+        kps0 = self.default_pts()
+        kps1 = warp_points_torch(kps0, data["H_0to1"], inverse=False)
+        pred = self.default_pred(kps0, kps1)
+
+        results = eval_matches_homography(data, pred)
+
+        self.assertEqual(results["prec@1px"], 1)
+        self.assertEqual(results["prec@3px"], 1)
+
+    def test_eval_matches_homography_real_outliers(self):
+        data = {"H_0to1": torch.tensor([[1.5, 0.2, 21], [-0.3, 1.6, 33], [0, 0, 1.0]])}
+        kps0 = self.default_pts()
+        kps0 = torch.cat([kps0, torch.tensor([[5.0, 5.0]])])
+        kps1 = warp_points_torch(kps0, data["H_0to1"], inverse=False)
+        # Move one keypoint 1.5 pixels away in x and y
+        kps1[-1] += 1.5
+        pred = self.default_pred(kps0, kps1)
+
+        results = eval_matches_homography(data, pred)
+        self.assertAlmostEqual(results["prec@1px"], 0.8)
+        self.assertAlmostEqual(results["prec@3px"], 1.0)
+
+    def test_eval_matches_homography_batched(self):
+        H0 = torch.tensor([[1.5, 0.2, 21], [-0.3, 1.6, 33], [0, 0, 1.0]])
+        H1 = torch.tensor([[0.7, 0.1, -5], [-0.1, 0.65, 13], [0, 0, 1.0]])
+        data = {"H_0to1": torch.stack([H0, H1])}
+        kps0 = torch.stack([self.default_pts(), self.default_pts().flip(0)])
+        kps1 = warp_points_torch(kps0, data["H_0to1"], inverse=False)
+        # In the first element of the batch there is one outlier
+        kps1[0, -1] += 5
+        matches0 = torch.stack([torch.arange(4), torch.arange(4)])
+        # In the second element of the batch there is only 2 matches
+        matches0[1, :2] = -1
+        pred = {
+            "keypoints0": kps0,
+            "keypoints1": kps1,
+            "matches0": matches0,
+            "matching_scores0": torch.ones_like(matches0),
+        }
+
+        results = eval_matches_homography(data, pred)
+        self.assertAlmostEqual(results["prec@1px"][0], 0.75)
+        self.assertAlmostEqual(results["prec@1px"][1], 1.0)
+        self.assertAlmostEqual(results["num_matches"][0], 4)
+        self.assertAlmostEqual(results["num_matches"][1], 2)

vismatch/third_party/keypt2subpx/submodules/glue_factory/tests/test_integration.py

@@ -0,0 +1,132 @@
+import unittest
+from collections import namedtuple
+from os.path import splitext
+
+import cv2
+import matplotlib.pyplot as plt
+import torch.cuda
+from kornia import image_to_tensor
+from omegaconf import OmegaConf
+from parameterized import parameterized
+from torch import Tensor
+
+from gluefactory import logger
+from gluefactory.eval.utils import (
+    eval_homography_dlt,
+    eval_homography_robust,
+    eval_matches_homography,
+)
+from gluefactory.models.two_view_pipeline import TwoViewPipeline
+from gluefactory.settings import root
+from gluefactory.utils.image import ImagePreprocessor
+from gluefactory.utils.tensor import map_tensor
+from gluefactory.utils.tools import set_seed
+from gluefactory.visualization.viz2d import (
+    plot_color_line_matches,
+    plot_images,
+    plot_matches,
+)
+
+
+def create_input_data(cv_img0, cv_img1, device):
+    img0 = image_to_tensor(cv_img0).float() / 255
+    img1 = image_to_tensor(cv_img1).float() / 255
+    ip = ImagePreprocessor({})
+    data = {"view0": ip(img0), "view1": ip(img1)}
+    data = map_tensor(
+        data,
+        lambda t: t[None].to(device)
+        if isinstance(t, Tensor)
+        else torch.from_numpy(t)[None].to(device),
+    )
+    return data
+
+
+ExpectedResults = namedtuple("ExpectedResults", ("num_matches", "prec3px", "h_error"))
+
+
+class TestIntegration(unittest.TestCase):
+    methods_to_test = [
+        ("superpoint+NN.yaml", "poselib", ExpectedResults(1300, 0.8, 1.0)),
+        ("superpoint-open+NN.yaml", "poselib", ExpectedResults(1300, 0.8, 1.0)),
+        (
+            "superpoint+lsd+gluestick.yaml",
+            "homography_est",
+            ExpectedResults(1300, 0.8, 1.0),
+        ),
+        (
+            "superpoint+lightglue-official.yaml",
+            "poselib",
+            ExpectedResults(1300, 0.8, 1.0),
+        ),
+    ]
+
+    visualize = False
+
+    @parameterized.expand(methods_to_test)
+    @torch.no_grad()
+    def test_real_homography(self, conf_file, estimator, exp_results):
+        set_seed(0)
+        model_path = root / "gluefactory" / "configs" / conf_file
+        img_path0 = root / "assets" / "boat1.png"
+        img_path1 = root / "assets" / "boat2.png"
+        h_gt = torch.tensor(
+            [
+                [0.85799, 0.21669, 9.4839],
+                [-0.21177, 0.85855, 130.48],
+                [1.5015e-06, 9.2033e-07, 1],
+            ]
+        )
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        gs = TwoViewPipeline(OmegaConf.load(model_path).model).to(device).eval()
+
+        cv_img0, cv_img1 = cv2.imread(str(img_path0)), cv2.imread(str(img_path1))
+        data = create_input_data(cv_img0, cv_img1, device)
+        pred = gs(data)
+        pred = map_tensor(
+            pred, lambda t: torch.squeeze(t, dim=0) if isinstance(t, Tensor) else t
+        )
+        data["H_0to1"] = h_gt.to(device)
+        data["H_1to0"] = torch.linalg.inv(h_gt).to(device)
+
+        results = eval_matches_homography(data, pred)
+        results = {**results, **eval_homography_dlt(data, pred)}
+
+        results = {
+            **results,
+            **eval_homography_robust(
+                data,
+                pred,
+                {"estimator": estimator},
+            ),
+        }
+
+        logger.info(results)
+        self.assertGreater(results["num_matches"], exp_results.num_matches)
+        self.assertGreater(results["prec@3px"], exp_results.prec3px)
+        self.assertLess(results["H_error_ransac"], exp_results.h_error)
+
+        if self.visualize:
+            pred = map_tensor(
+                pred, lambda t: t.cpu().numpy() if isinstance(t, Tensor) else t
+            )
+            kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+            m0 = pred["matches0"]
+            valid0 = m0 != -1
+            kpm0, kpm1 = kp0[valid0], kp1[m0[valid0]]
+
+            plot_images([cv_img0, cv_img1])
+            plot_matches(kpm0, kpm1, a=0.0)
+            plt.savefig(f"{splitext(conf_file)[0]}_point_matches.svg")
+
+            if "lines0" in pred and "lines1" in pred:
+                lines0, lines1 = pred["lines0"], pred["lines1"]
+                lm0 = pred["line_matches0"]
+                lvalid0 = lm0 != -1
+                linem0, linem1 = lines0[lvalid0], lines1[lm0[lvalid0]]
+
+                plot_images([cv_img0, cv_img1])
+                plot_color_line_matches([linem0, linem1])
+                plt.savefig(f"{splitext(conf_file)[0]}_line_matches.svg")
+            plt.show()

vismatch/third_party/keypt2subpx/summarize.py

@@ -0,0 +1,44 @@
+import argparse
+import pickle
+import numpy as np
+import glob
+
+from utils import pose_auc
+
+parser = argparse.ArgumentParser(
+    description="Summarize Test Set Results.",
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+parser.add_argument('experiment', type=str, help='experiment name')
+
+opt, unparsed = parser.parse_known_args()
+if len(unparsed) > 0:
+    parser.print_usage()
+    exit(1)
+
+filelist = glob.glob(f"results/{opt.experiment}/pickled_results/*.pkl")
+
+err_t, err_r, inliers, outliers = [], [], [], []
+for filename in filelist:
+    with open(filename, 'rb') as f:
+        data = pickle.load(f)
+        err_t += data['error_t']
+        err_r += data['error_R']
+        inliers += data['inliers']
+        outliers += data['outliers']
+
+pose_errors = []
+for idx in range(len(err_t)):
+    pose_error = np.maximum(err_t[idx], err_r[idx])
+    if pose_error == np.inf:
+        pose_error = 180
+    pose_errors.append(pose_error)
+
+thresholds = [5, 10, 20]
+aucs = pose_auc(pose_errors, thresholds)
+aucs = [100.*yy for yy in aucs]
+
+print('Evaluation Results (mean over {} pairs):'.format(len(err_t)))
+print('AUC@5\t AUC@10\t AUC@20\tInRat.\tMean\tMedian\t')
+print('{:.2f}\t {:.2f}\t {:.2f}\t {:.2f}\t {:.2f}\t {:.2f}'.format(
+    aucs[0], aucs[1], aucs[2], 100.*np.sum(inliers) / (np.sum(inliers) + np.sum(outliers)),
+    np.mean(pose_errors), np.median(pose_errors)))

vismatch/third_party/keypt2subpx/test.py

@@ -0,0 +1,225 @@
+import numpy as np
+import cv2
+
+import torch
+import pickle
+from tqdm import tqdm
+from pathlib import Path
+
+import pygcransac
+
+from model import AttnTuner
+from utils import compute_pose_error, pose_auc, calculate_loss, create_log_dir
+from logger import Logger
+from dataset import SparsePatchDataset
+from settings import get_config, get_logger, print_usage
+
+CUDA_LAUNCH_BLOCKING=2, 3
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+
+def valid(valid_loader, model, opt, writer=None, step=-1):
+    model.eval()
+    err_ts, err_Rs = [], []
+    loss_vals = []
+    inlier_cnt = []
+    outlier_cnt = []
+
+    with torch.no_grad():
+        for idx, valid_data in enumerate(tqdm(valid_loader)):
+            corr, gt_F, gt_E, gt_R, gt_t, K1s, K2s, im_size1, im_size2, patch1, patch2, descriptor1, descriptor2 = valid_data[:13]
+            scorepatch1, scorepatch2 = (valid_data[13].cuda(), valid_data[14].cuda()) if len(valid_data) > 13 else (None, None)
+
+            ret = None
+            if len(corr.shape) > 1 and corr.shape[1] > 4:
+
+                if opt.vanilla: # without our model
+                    inp = corr[...,:4].cuda() # B x N x 4
+                else: # with our model
+                    meanft = ((descriptor1 + descriptor2) / 2.).cuda()
+                    idx1 = 2.5* K1s[:,:2,:2].unsqueeze(1).cuda().inverse() @ model(patch1.cuda(), scorepatch1, meanft).unsqueeze(-1)
+                    idx1 = idx1.squeeze(-1)
+                    idx2 = 2.5* K2s[:,:2,:2].unsqueeze(1).cuda().inverse() @ model(patch2.cuda(), scorepatch2, meanft).unsqueeze(-1)
+                    idx2 = idx2.squeeze(-1)
+                    inp = corr[...,:4].cuda() + torch.cat([idx1, idx2], 2) # B x N x 4
+
+                K1s_cuda = K1s.cuda()
+                K2s_cuda = K2s.cuda()
+
+                kpts0 = torch.cat([inp.squeeze()[:, :2], torch.ones_like(inp.squeeze()[:, 0:1])], dim=-1)
+                kpts1 = torch.cat([inp.squeeze()[:, 2:], torch.ones_like(inp.squeeze()[:, 0:1])], dim=-1)
+                # Let's change mkpts' normalized points to image coordinates
+
+                mkpts0 = kpts0@K1s_cuda.transpose(-1, -2).squeeze(0)
+                mkpts1 = kpts1@K2s_cuda.transpose(-1, -2).squeeze(0)
+
+                mkpts0 = mkpts0[:, :2] / mkpts0[:, 2:]
+                mkpts1 = mkpts1[:, :2] / mkpts1[:, 2:]
+
+                mkpts = torch.cat([mkpts0, mkpts1], dim=-1).cpu().detach().numpy()
+
+                # Let's change mkpts' normalized points to image coordinates
+                im_size1 = im_size1.detach().numpy()[0]
+                im_size2 = im_size2.detach().numpy()[0]
+                K1 = K1s.detach().numpy()[0]
+                K2 = K2s.detach().numpy()[0]
+                E, mask = pygcransac.findEssentialMatrix(
+                    np.ascontiguousarray(mkpts), 
+                    K1, K2,
+                    im_size1[0], im_size1[1], im_size2[0], im_size2[1],
+                    probabilities = [],
+                    threshold = opt.ransac_thr,
+                    conf = 0.99999, # RANSAC confidence
+                    max_iters = 1000,
+                    min_iters = 1000,
+                    sampler = 0)
+
+                kpts0 = kpts0[:, :2].cpu().detach().numpy()
+                kpts1 = kpts1[:, :2].cpu().detach().numpy()
+
+                mask = np.expand_dims(mask, axis=-1).astype(np.uint8)
+
+                best_num_inliers = 0
+                if E is not None:
+                    for _E in np.split(E, len(E) / 3):
+                        n, R, t, mask_ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask)
+                        if n > best_num_inliers:
+                            best_num_inliers = n
+                            ret = (R, t[:, 0], mask.ravel() > 0)
+
+                loss = calculate_loss(inp, gt_E.cuda(), K1s_cuda, K2s_cuda, opt.train_thr)
+                tmp = loss.cpu().detach().item()
+                if np.isnan(tmp):
+                    tqdm.write(f"Pair #{idx}: Loss is nan!!")
+                else:
+                    loss_vals.append(tmp)
+                
+            if ret is None:
+                tqdm.write(f"Pair #{idx}: Not enough points for E estimation or just E estimation failed")
+                err_t, err_R = np.inf, np.inf
+                inlier_cnt.append(0)
+                outlier_cnt.append(corr.shape[1])
+            else:
+                R, t, inliers = ret
+                T_0to1 = torch.cat([gt_R.squeeze(), gt_t.squeeze().unsqueeze(-1)], dim=-1).numpy()
+                err_t, err_R = compute_pose_error(T_0to1, R, t)
+                n = np.sum(inliers)
+                outlier_cnt.append(len(mkpts) - n)
+                inlier_cnt.append(n)
+
+            err_ts.append(err_t)
+            err_Rs.append(err_R)
+
+    # Write the evaluation results to disk.
+    out_eval = {
+        'error_t': err_ts,
+        'error_R': err_Rs,
+        'inliers': inlier_cnt,
+        'outliers': outlier_cnt
+    }
+
+    if opt.test:
+        pickle.dump(out_eval, open(opt.res_path / f'results_{opt.current_split}_{opt.total_split}.pkl', 'wb'))
+
+    pose_errors = []
+    for idx in range(len(out_eval['error_t'])):
+        pose_error = np.maximum(out_eval['error_t'][idx], out_eval['error_R'][idx])
+        if pose_error == np.inf:
+            pose_error = 180
+        pose_errors.append(pose_error)
+
+    thresholds = [5, 10, 20]
+    aucs = pose_auc(pose_errors, thresholds)
+    aucs = [100.*yy for yy in aucs]
+    in_rate = 100. * np.sum(inlier_cnt) / (np.sum(inlier_cnt) + np.sum(outlier_cnt))
+    err_mean = np.mean(pose_errors)
+    err_median = np.median(pose_errors)
+    loss = np.mean(loss_vals)
+
+    if writer is not None:
+        writer.add_scalar('val/loss', loss, step)
+        writer.add_scalar('val/inlier_ratio', np.sum(inlier_cnt) / (np.sum(inlier_cnt) + np.sum(outlier_cnt)), step)
+        writer.add_scalar('val/AUC@5', aucs[0], step)
+        writer.add_scalar('val/AUC@10', aucs[1], step)
+        writer.add_scalar('val/AUC@20', aucs[2], step)
+
+    print(f'Evaluation Results (mean over {len(err_ts)} pairs):')
+    print('AUC@5\t AUC@10\t AUC@20\t InRat.\t Mean\t Median\t Avg loss\t')
+    print(f'{aucs[0]:.2f}\t {aucs[1]:.2f}\t {aucs[2]:.2f}\t '
+          f'{in_rate:.2f}\t {err_mean:.2f}\t {err_median:.2f}\t {loss:.8f}')
+
+    return aucs[0], aucs[1], aucs[2], in_rate, err_mean, err_median, loss
+
+if __name__ == "__main__":
+    # parse command line arguments
+    # If we have unparsed arguments, print usage and exit
+    opt, unparsed = get_config()
+    if len(unparsed) > 0:
+        print_usage()
+        exit(1)
+    logger = get_logger(opt)
+
+    # construct folder that should contain pre-calculated correspondences
+    test_data = opt.datasets.split(',') #support multiple training datasets used jointly
+    logger.info('Using datasets: [' + ', '.join(test_data) + ']')
+
+    # create or load model
+    lentable = {
+        'spnn': 256,
+        'splg': 256,
+        'aliked': 128,
+        'dedode': 256,
+        'xfeat': 64
+    }
+    with_score = opt.detector in ['spnn', 'splg', 'aliked']
+    model = AttnTuner(lentable[opt.detector], with_score)
+    if len(opt.model) > 0:
+        model.load_state_dict(torch.load(opt.model)['model'])
+    model = model.cuda()
+    model.eval()
+
+    # ----------------------------------------
+    logger.info(f"Starting experiment {opt.experiment}")
+    output_dir = Path(__file__).parent / 'results' / opt.experiment
+    output_dir.mkdir(exist_ok=True, parents=True)
+    create_log_dir(output_dir, opt)
+
+    if opt.detect_anomaly:
+        torch.autograd.set_detect_anomaly(True)
+
+    if opt.current_split > 1 and opt.total_split > 1:
+        print(f"Working with split: {opt.current_split} / {opt.total_split}", flush=True)
+
+    filelogger = Logger(opt.log_path / 'log_evaluation.txt', title='k2s')
+    filelogger.set_names(['AUC5'] + ['AUC10', 'AUC20', 'InRat.', 'Mean', 'Median', 'Loss'])
+    
+    if opt.test: # Metrics: aucs5, aucs10, aucs20, in_rat, err_mean, err_median, loss_avg
+        test_data = ['data/' + ds + '/' + 'test_' + opt.detector + '/' for ds in test_data]
+        testset = SparsePatchDataset(test_data, opt.nfeatures, opt.fmat, not opt.sideinfo, not with_score,#False,
+                current_split=opt.current_split, total_split=opt.total_split)
+        test_loader = torch.utils.data.DataLoader(testset, shuffle=False, num_workers=6, batch_size=1)
+        aucs5, aucs10, aucs20, in_rat, err_mean, err_median, loss_avg = valid(test_loader, model, opt)
+        filelogger.append([aucs5, aucs10, aucs20, in_rat, err_mean, err_median, loss_avg])
+    else:
+        val_data = ['data/' + ds + '/' + 'val_' + opt.detector + '/' for ds in test_data]
+        valset = SparsePatchDataset(val_data, opt.nfeatures, opt.fmat, not opt.sideinfo, not with_score,#False,
+                current_split=opt.current_split, total_split=opt.total_split)
+        val_loader = torch.utils.data.DataLoader(valset, shuffle=False, num_workers=6, batch_size=1)
+        all_metrics = None
+        for i in range(opt.total_run):
+            logger.info(f"Test run {i+1}/{opt.total_run}")
+            metrics = valid(val_loader, model, opt)
+            filelogger.append(metrics)
+            if all_metrics is None:
+                all_metrics = [[m] for m in metrics]
+            else:
+                for j in range(len(all_metrics)):
+                    all_metrics[j].append(metrics[j])
+        all_metrics = [np.mean(m) for m in all_metrics]
+
+        print('----------------------------------------')
+        print(f'Final Evaluation Results (average over {opt.total_run} runs):')
+        print('AUC@5\t AUC@10\t AUC@20\t InRat.\t Mean\t Median\t Avg loss\t')
+        print(f'{all_metrics[0]:.2f}\t {all_metrics[1]:.2f}\t {all_metrics[2]:.2f}\t '
+            f'{all_metrics[3]:.2f}\t {all_metrics[4]:.2f}\t {all_metrics[5]:.2f}\t {all_metrics[6]:.8f}')

vismatch/third_party/keypt2subpx/train.py

@@ -0,0 +1,180 @@
+import os
+
+import torch
+import torch.optim as optim
+from tqdm import trange
+from tensorboardX import SummaryWriter
+from pathlib import Path
+
+from model import AttnTuner
+from utils import calculate_loss, create_log_dir, last_checkpoint
+from test import valid
+from logger import Logger
+from dataset import SparsePatchDataset
+from settings import get_config, get_logger, print_usage
+
+CUDA_LAUNCH_BLOCKING=2, 3
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+
+def train_step(step, optimizer, model, data, writer, mask=None):
+    corr, _, gt_E, _, _, K1s, K2s, _, _, patch1, patch2, descriptor1, descriptor2 = data[:13]
+    scorepatch1, scorepatch2 = (data[13].cuda(), data[14].cuda()) if len(data) > 13 else (None, None)
+    meanft = ((descriptor1 + descriptor2) / 2.).cuda()
+    idx1 = 2.5 * K1s[:,:2,:2].unsqueeze(1).cuda().inverse() @ model(patch1.cuda(), scorepatch1, meanft).unsqueeze(-1)
+    idx1 = idx1.squeeze(-1)
+    idx2 = 2.5 * K2s[:,:2,:2].unsqueeze(1).cuda().inverse() @ model(patch2.cuda(), scorepatch2, meanft).unsqueeze(-1)
+    idx2 = idx2.squeeze(-1)
+    inp = corr[...,:4].cuda() + torch.cat([idx1, idx2], 2) # B x N x 4
+    loss = calculate_loss(inp, gt_E.cuda(), K1s.cuda(), K2s.cuda(), opt.train_thr)
+        
+    writer.add_scalar('train/loss', loss.item(), step+1)
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    return [loss.item()], mask
+
+def train(model, train_loader, valid_loader, opt, writer, logger):
+    optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)
+
+    checkpoint_path = last_checkpoint(opt.model_path)
+    opt.resume = opt.resume and bool(checkpoint_path) and os.path.isfile(checkpoint_path)
+    if opt.resume:
+        logger.info('==> Resuming from checkpoint..')
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+        best_acc = checkpoint['best_acc']
+        start_step = checkpoint['step']
+        model.load_state_dict(checkpoint['model'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+
+        logger_train = Logger(opt.log_path / 'log_train.txt', title='k2s', resume=True)
+        logger_valid = Logger(opt.log_path / 'log_valid.txt', title='k2s', resume=True)
+
+    else:
+        start_step = 0
+
+        logger_train = Logger(opt.log_path / 'log_train.txt', title='k2s')
+        logger_train.set_names(['Learning Rate'] + ['Sampson Loss'])
+        logger_valid = Logger(opt.log_path / 'log_valid.txt', title='k2s')
+        logger_valid.set_names(['AUC5'] + ['AUC10', 'AUC20'])
+
+        aucs5, aucs10, aucs20, _, _, _, loss = valid(valid_loader, model, opt, writer, start_step)
+        logger_valid.append([aucs5, aucs10, aucs20])
+
+        best_acc = -loss
+        logger.info(f"Saving initial model with va_res = {best_acc:6.3f}")
+        if start_step == 0:
+            torch.save({
+                'step': start_step,
+                'model': model.state_dict(),
+                'best_acc': best_acc,
+                'optimizer' : optimizer.state_dict(),
+            }, os.path.join(opt.model_path, f'checkpoint_{start_step}.pth'))
+
+    train_loader_iter = iter(train_loader)
+
+    tbar = trange(start_step, opt.train_iter)
+
+    mask = None
+
+    for step in tbar:
+        try:
+            train_data = next(train_loader_iter)
+        except StopIteration:
+            train_loader_iter = iter(train_loader)
+            train_data = next(train_loader_iter)
+
+        # run training
+        cur_lr = optimizer.param_groups[0]['lr']
+
+        loss_vals, mask = train_step(step, optimizer, model, train_data, writer, mask)
+
+        tbar.set_description('Doing: {}/{}, LR: {}, Sampson_loss: {}'\
+        .format(step+1, opt.train_iter, cur_lr, loss_vals[0]))
+
+        if step % 100 == 0:
+            logger_train.append([cur_lr] + loss_vals)
+
+        # Check if we want to write validation
+        b_save = ((step+1) % opt.save_intv) == 0
+        b_validate = ((step+1) % opt.val_intv) == 0
+
+        if b_validate:
+            aucs5, aucs10, aucs20, _, _, _, loss = valid(valid_loader, model, opt, writer, step)
+            logger_valid.append([aucs5, aucs10, aucs20])
+
+            va_res = -loss
+            if va_res > best_acc:
+                logger.info(f"Saving best model with va_res = {best_acc:6.3f}")
+                best_acc = va_res
+                torch.save({
+                'step': step + 1,
+                'model': model.state_dict(),
+                'best_acc': best_acc,
+                'optimizer' : optimizer.state_dict(),
+                }, os.path.join(opt.model_path, 'model_best.pth'))
+
+            model.train()
+
+        if b_save:
+            torch.save({
+            'step': step + 1,
+            'model': model.state_dict(),
+            'best_acc': best_acc,
+            'optimizer' : optimizer.state_dict(),
+            }, os.path.join(opt.model_path, f'checkpoint_{step+1}.pth'))
+
+if __name__ == "__main__":
+    # parse command line arguments
+    # If we have unparsed arguments, print usage and exit
+    opt, unparsed = get_config()
+    if len(unparsed) > 0:
+        print_usage()
+        exit(1)
+    logger = get_logger(opt)
+
+    # construct folder that should contain pre-calculated correspondences
+    train_data = opt.datasets.split(',') #support multiple training datasets used jointly
+    logger.info('Using datasets: ' + ', '.join(train_data))
+
+    valid_data = train_data
+    train_data = ['data/' + ds + '/' + 'train_' + opt.detector + '/' for ds in train_data]
+    valid_data = ['data/' + ds + '/' + 'val_' + opt.detector + '/' for ds in valid_data]
+
+    with_score = opt.detector in ['spnn', 'splg', 'aliked']
+
+    trainset = SparsePatchDataset(train_data, opt.nfeatures, opt.fmat, not opt.sideinfo, not with_score, train=True)
+    valset = SparsePatchDataset(valid_data, opt.nfeatures, opt.fmat, not opt.sideinfo, not with_score)
+    train_loader = torch.utils.data.DataLoader(trainset, shuffle=True, num_workers=8, batch_size=opt.batchsize)
+    valid_loader = torch.utils.data.DataLoader(valset, shuffle=False, num_workers=8, batch_size=1)
+
+    logger.info(f"Image pairs: {len(trainset)}")
+
+    # create or load model
+    lentable = {
+        'spnn': 256,
+        'splg': 256,
+        'aliked': 128,
+        'dedode': 256,
+        'xfeat': 64
+    }
+    model = AttnTuner(lentable[opt.detector], with_score)
+    if len(opt.model) > 0:
+        model.load_state_dict(torch.load(opt.model)['model'])
+    model = model.cuda()
+    model.train()
+
+    # ----------------------------------------
+    logger.info(f"Starting experiment {opt.experiment}")
+    output_dir = Path(__file__).parent / 'results' / opt.experiment
+    output_dir.mkdir(exist_ok=True, parents=True)
+    create_log_dir(output_dir, opt)
+    writer = SummaryWriter(log_dir=str(output_dir))
+
+    if opt.detect_anomaly:
+        torch.autograd.set_detect_anomaly(True)
+
+    train(model, train_loader, valid_loader, opt, writer, logger)