import torch
import torchvision.transforms as tvf

from ..utils.base_model import BaseModel


class OpenIBL(BaseModel):
    default_conf = {
        'model_name': 'vgg16_netvlad',
    }
    required_inputs = ['image']

    def _init(self, conf):
        self.net = torch.hub.load(
            'yxgeee/OpenIBL', conf['model_name'], pretrained=True).eval()
        mean = [0.48501960784313836, 0.4579568627450961, 0.4076039215686255]
        std = [0.00392156862745098, 0.00392156862745098, 0.00392156862745098]
        self.norm_rgb = tvf.Normalize(mean=mean, std=std)

    def _forward(self, data):
        image = self.norm_rgb(data['image'])
        desc = self.net(image)
        return {
            'global_descriptor': desc,
        }
