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

__all__ = ['optimize_points_H', 'optimize_points_H_hardnet']

# Cell
import torch
import kornia as K
import kornia.geometry as KG
import kornia.feature as KF
import cv2
import numpy as np

from copy import deepcopy

def optimize_points_H(img1, img2, pts1, pts2, n_iter=1):
    min_size = min(pts1.shape[0], pts2.shape[0])

    def get_bbox(pts_):
        pts = np.array(pts_).astype(np.int32)
        min_x = pts[:,0].min()
        min_y = pts[:,1].min()
        max_x = pts[:,0].max()
        max_y = pts[:,1].max()
        return min_x,min_y,max_x,max_y
    num = min_size
    min_x1,min_y1,max_x1,max_y1  = get_bbox(pts1)
    min_x2,min_y2,max_x2,max_y2  = get_bbox(pts2)
    pts1_centered = np.array(pts1[:min_size]) - np.array([min_x1, min_y1])
    pts2_centered = np.array(pts2[:min_size]) - np.array([min_x2, min_y2])
    timg1 = K.image_to_tensor(img1[min_y1:max_y1, min_x1:max_x1], None).float()#.mean(dim=1, keepdim=True)
    timg2 = K.image_to_tensor(img2[min_y2:max_y2, min_x2:max_x2], None).float()#.mean(dim=1, keepdim=True)

    pts1_opt = deepcopy(torch.from_numpy(pts1_centered)).float().reshape(1, -1, 2).requires_grad_()
    pts2_opt = deepcopy(torch.from_numpy(pts2_centered)).float().reshape(1, -1, 2).requires_grad_()
    #opt = torch.optim.Adam((pts1_opt, pts2_opt), lr=1.0)
    opt = torch.optim.LBFGS((pts1_opt, pts2_opt), lr=0.1,max_iter=20)
    desc = KF.HardNet(True).eval()
    sp2 = KG.build_pyramid(timg2, 3)
    hull2 = cv2.convexHull(pts2_centered.astype(np.float32))
    img2_c = img2[min_y2:max_y2, min_x2:max_x2]
    mask = K.image_to_tensor(cv2.fillPoly(img2_c,
                        hull2.astype(np.int32).reshape(1,-1,2),
                        color=(255, 255, 255)), None).float()
    mask_pyr = KG.build_pyramid(mask, 3)
    for i in range(n_iter):
        def closure():
            opt.zero_grad()
            H = KG.find_homography_dlt(pts1_opt, pts2_opt)
            img1_warped = K.warp_perspective(timg1, H, timg2.shape[2:])
            sp1 = KG.build_pyramid(img1_warped, 3)
            loss = torch.nn.functional.l1_loss(sp1[2]*(mask_pyr[2]>0).float(),#[:,:,300:600],
                                               sp2[2]*(mask_pyr[2]>0).float())
            loss += torch.nn.functional.l1_loss(sp1[1]*(mask_pyr[1]>0).float(),#[:,:,300:600],
                                               sp2[1]*(mask_pyr[1]>0).float())
            #[:,:,300:600])# +  torch.nn.functional.mse_loss(sp1[2], sp2[2])+ torch.nn.functional.mse_loss(sp1[0], sp2[0])
            loss.backward()
            print (loss.item())
            return loss

        #lafs2 = KF.laf_from_center_scale_ori(pts2_opt,
        #                                     32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
        #                                     10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
        #pts1_reproject = KG.transform_points(H, pts1_opt)
        #lafs1_in_2 = KF.laf_from_center_scale_ori(pts1_reproject,
        #                                     32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
        #                                     10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
        #patches1_in2 = KF.extract_patches_from_pyramid(img1_warped, lafs1_in_2, 32)[0]
        #patches2 = KF.extract_patches_from_pyramid(timg2, lafs2, 32)[0]
        #plt.figure()
        #plt.imshow(K.tensor_to_image(patches1_in2[4]-patches2[4]))
        #plt.imshow(K.tensor_to_image(img1_warped))

        #patches1_in2 = desc(patches1_in2)
        #patches2 = desc(patches2)
        #print (img1_warped.shape, timg2.shape)
        #loss = torch.nn.functional.mse_loss(img1_warped, timg2)
        #opt.zero_grad()
        #loss.backward()
        opt.step(closure)
    return pts1_opt+torch.tensor([min_x1, min_y1]), pts2_opt+torch.tensor([min_x2, min_y2])


# Cell
import torch
import kornia as K
import kornia.geometry as KG
import kornia.feature as KF
import cv2
import numpy as np

from copy import deepcopy

def optimize_points_H_hardnet(img1, img2, pts1, pts2, n_iter=1):
    timg1 = K.image_to_tensor(img1, None).float().mean(dim=1, keepdim=True)
    timg2 = K.image_to_tensor(img2, None).float().mean(dim=1, keepdim=True)
    min_size = min(pts1.shape[0], pts2.shape[0])
    pts1_opt = deepcopy(torch.from_numpy(pts1)).float().reshape(1, -1, 2).requires_grad_()
    pts2_opt = deepcopy(torch.from_numpy(pts2)).float().reshape(1, -1, 2).requires_grad_()
    #opt = torch.optim.Adam((pts1_opt, pts2_opt), lr=1.0)
    opt = torch.optim.LBFGS((pts1_opt, pts2_opt), lr=2.0,max_iter=10)
    desc = KF.HardNet(True).eval()
    sp2 = KG.build_pyramid(timg2, 3)
    for i in range(n_iter):
        def closure():
            opt.zero_grad()
            H1 = KG.find_homography_dlt(pts1_opt, pts2_opt)
            img1_warped = K.warp_perspective(timg1, H1, img2.shape[:2])
            #sp1 = KG.build_pyramid(img1_warped, 3)
            #loss = torch.nn.functional.l1_loss(sp1[2][:,:,300:600],
            #                                   sp2[2][:,:,300:600])# +  torch.nn.functional.mse_loss(sp1[2], sp2[2])+ torch.nn.functional.mse_loss(sp1[0], sp2[0])
            lafs2 = KF.laf_from_center_scale_ori(pts2_opt,
                                         32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
                                             10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
            pts1_reproject = KG.transform_points(H1, pts1_opt)
            lafs1_in_2 = KF.laf_from_center_scale_ori(pts1_reproject,
                                                 32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
                                                 10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
            patches1_in2 = KF.extract_patches_from_pyramid(img1_warped, lafs1_in_2, 32)[0]
            patches2 = KF.extract_patches_from_pyramid(timg2, lafs2, 32)[0]
            #plt.figure()
            #plt.imshow(K.tensor_to_image(patches1_in2[4]-patches2[4]))
            #plt.imshow(K.tensor_to_image(img1_warped))

            patches1_in2 = desc(patches1_in2)
            patches2 = desc(patches2)
            loss = (patches1_in2 - patches2).pow(2).sum(dim=-1).sqrt().mean()
            loss.backward()
            print (loss.item())
            return loss

        #lafs2 = KF.laf_from_center_scale_ori(pts2_opt,
        #                                 32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
        #                                     10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
        #pts1_reproject = KG.transform_points(H, pts1_opt)
        #lafs1_in_2 = KF.laf_from_center_scale_ori(pts1_reproject,
        #                                     32*torch.ones(1, min_size, 1, 1, dtype=pts2_opt.dtype),
        #                                     10*torch.zeros(1, min_size, 1, dtype=pts2_opt.dtype))
        #patches1_in2 = KF.extract_patches_from_pyramid(img1_warped, lafs1_in_2, 32)[0]
        #patches2 = KF.extract_patches_from_pyramid(timg2, lafs2, 32)[0]
        #plt.figure()
        #plt.imshow(K.tensor_to_image(patches1_in2[4]-patches2[4]))
        #plt.imshow(K.tensor_to_image(img1_warped))

        #patches1_in2 = desc(patches1_in2)
        #patches2 = desc(patches2)
        #print (img1_warped.shape, timg2.shape)
        #loss = torch.nn.functional.mse_loss(img1_warped, timg2)
        #opt.zero_grad()
        #loss.backward()
        opt.step(closure)
    return pts1_opt, pts2_opt
