import sys
from pathlib import Path

from unav.third_party.local_feature.SuperPoint_SuperGlue.base_model import BaseModel

sys.path.append(str(Path(__file__).parent))
from unav.third_party.local_feature.SuperPoint_SuperGlue.SuperGluePretrainedNetwork.models.superpoint import SuperPoint as SP


class SuperPoint(BaseModel):
    default_conf = {
        'nms_radius': 4,
        'keypoint_threshold': 0.005,
        'max_keypoints': -1,
        'remove_borders': 4,
    }
    required_inputs = ['image']

    def _init(self, conf):
        self.net = SP(conf)

    def _forward(self, data):
        return self.net(data)
