jit.py

import torch

from models.superglue_triton import SuperGlue
from models.superpoint_triton import SuperPoint
from models.superpoint_descriptor import SuperPointDescriptor


def export_sp():
    superpoint = SuperPoint({
        'max_keypoints': 5000,
        'keypoint_threshold': 0.01
    }).eval()
    torch.jit.save(superpoint, 'superpoint_model.pt')


def export_sg():
    superglue = SuperGlue({
        'weights': 'outdoor',
        'sinkhorn_iterations': 10
    }).eval()
    torch.jit.save(superglue, 'superglue_model.pt')


def export_superpoint_desc():
    superpoint_desc = SuperPointDescriptor().eval()
    torch.jit.save(superpoint_desc, 'superpoint_desc_model.pt')


export_sp()
# export_sg()
# export_superpoint_desc()