# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/facenet.ipynb.

# %% auto 0
__all__ = ['FaceNetInceptionResnetV1', 'facenet']

# %% ../nbs/facenet.ipynb 3
from facenet_pytorch import InceptionResnetV1, fixed_image_standardization

from .utils import *
from .siamese import *
from .pair_matching import *


class FaceNetInceptionResnetV1(InceptionResnetV1):
    """Like `pytorch_facenet`'s `InceptionResnetV1` but accepts standard float tensors"""
    def __init__(self, pretrained=True, classify=True):
        if pretrained is True:
            pretrained = 'vggface2'
        super().__init__(pretrained, classify)

    def forward(self, x):
        x = fixed_image_standardization(x * 255)
        return super().forward(x)


def facenet(pretrained=True) -> ThresholdSiamese:
    """A binary classifier matching pairs of facial images"""
    backbone = FaceNetInceptionResnetV1(pretrained, classify=False)
    backbone = cut_model_by_name(backbone, 'last_bn')  # skip normalization in case we want unnormalized features
    return ThresholdSiamese(backbone)
