# AUTOGENERATED! DO NOT EDIT! File to edit: deephelper.ipynb (unless otherwise specified).

__all__ = ['Bunch', 'COTRHelper', 'cotr_default_opts']

# Cell
import cv2
import os
import kornia as K
import kornia.feature as KF
import kornia.geometry as KG
import os
import torch
from copy import deepcopy
from kornia_moons.feature import *
from .io import *
from .metrics import *
from .core import *


import numpy as np
import torch
import imageio

from .COTR.COTR.utils import utils, debug_utils
from .COTR.COTR.models import build_model
from .COTR.COTR.options.options import *
from .COTR.COTR.options.options_utils import *
from .COTR.COTR.inference.inference_helper import triangulate_corr
from .COTR.COTR.inference.sparse_engine import SparseEngine

class Bunch(object):
    def __init__(self, adict):
        self.__dict__.update(adict)

cotr_default_opts = {'backbone': 'resnet50',
 'hidden_dim': 256,
 'dilation': False,
 'dropout': 0.1,
 'nheads': 8,
 'layer': 'layer3',
 'enc_layers': 6,
 'dec_layers': 6,
 'position_embedding': 'lin_sine',
 'out_dir': './out',
 'load_weights': None,
 'faster_infer': True,
 'command': '',
 'dim_feedforward': 1024,
 'load_weights_path': '/Users/oldufo/dev/correspondence_annotator/pixelstitch/COTR/out/default/checkpoint.pth.tar'}


class COTRHelper():
    def __init__(self, opt=cotr_default_opts, on_gpu = False):
        self.on_gpu = on_gpu
        opt = Bunch(opt)
        self.model = build_model(opt)
        weights = torch.load(opt.load_weights_path, map_location='cpu')['model_state_dict']
        utils.safe_load_weights(self.model, weights)
        if self.on_gpu:
            self.model = self.model.cuda()
        self.model = self.model.eval()
        self.engine = SparseEngine(self.model, 32, mode='tile')
        return
    def add_images(self, img_a, img_b):
        self.img_a = img_a
        self.img_b = img_b
        return
    def suggest_correspondence(self, xy, on_left: bool):
        num =len(xy)
        assert xy.shape[1] == 2
        img_a_area = 1.0
        img_b_area = 1.0
        corrs = self.engine.cotr_corr_multiscale(self.img_a if on_left else self.img_b,
                                                 self.img_b if on_left else self.img_a,
                                                 np.linspace(0.75, 0.1, 4), #np.linspace(0.5, 0.0625, 4),
                                                  1, max_corrs=num, queries_a=xy, force=True, areas=[img_a_area, img_b_area])
        return corrs