# AUTOGENERATED FROM: babilim/model/imagenet.ipynb

# Cell: 0
"""doc
# babilim.model.imagenet

> An implemntation of various imagenet models.
"""

# Cell: 1
from babilim.core.annotations import RunOnlyOnce
from babilim.core.module_native import ModuleNative


# Cell: 2
class ImagenetModel(ModuleNative):
    def __init__(self, encoder_type, only_encoder=False, pretrained=False, last_layer=None):
        """
        Create one of the iconic image net models in one line.
        Allows for only using the encoder part.

        This model assumes the input image to be 0-255 (8 bit integer) with 3 channels.

        :param encoder_type: The encoder type that should be used. Must be in ("vgg16", "vgg16_bn", "vgg19", "vgg19_bn", "resnet50", "resnet101", "resnet152", "densenet121", "densenet169", "densenet201", "inception_v3", "mobilenet_v2")
        :param only_encoder: Leaves out the classification head for VGG16 leaving you with a feature encoder.
        :param pretrained: If you want imagenet weights for this network.
        :param last_layer: Index of the last layer in the encoder. Allows to cutoff encoder after a few layers.
        """
        super().__init__()
        self.only_encoder = only_encoder
        self.pretrained = pretrained
        self.encoder_type = encoder_type
        self.last_layer = last_layer

    @RunOnlyOnce
    def _build_tf(self, image):
        raise NotImplementedError()

    def _call_tf(self, image):
        raise NotImplementedError()

    @RunOnlyOnce
    def _build_pytorch(self, image):
        import torch
        from torchvision.models import vgg16, vgg16_bn, vgg19, vgg19_bn, resnet50, resnet101, resnet152, densenet121, densenet169, densenet201, inception_v3, mobilenet_v2
        from torch.nn import Sequential
        model = None
        if self.encoder_type == "vgg16":
            model = vgg16
        elif self.encoder_type == "vgg16_bn":
            model = vgg16_bn
        elif self.encoder_type == "vgg19":
            model = vgg19
        elif self.encoder_type == "vgg19_bn":
            model = vgg19_bn
        elif self.encoder_type == "resnet50":
            model = resnet50
        elif self.encoder_type == "resnet101":
            model = resnet101
        elif self.encoder_type == "resnet152":
            model = resnet152
        elif self.encoder_type == "densenet121":
            model = densenet121
        elif self.encoder_type == "densenet169":
            model = densenet169
        elif self.encoder_type == "densenet201":
            model = densenet201
        elif self.encoder_type == "inception_v3":
            model = inception_v3
        elif self.encoder_type == "mobilenet_v2":
            model = mobilenet_v2
        else:
            raise RuntimeError("Unsupported encoder type.")
        
        if self.only_encoder:
            encoder = list(model(pretrained=self.pretrained).features)
            if self.last_layer is not None:
                encoder = encoder[:self.last_layer+1]
            self.model = Sequential(*encoder)
        else:
            self.model = model(pretrained=self.pretrained)
        
        if torch.cuda.is_available():
            self.model = self.model.to(torch.device(self.device))
        
        # Just in case, make the image a float tensor
        image = image.float()

        # Standardization values from torchvision.models documentation
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        # Create tensors for a 0-255 value range image.
        self.mean = torch.as_tensor([i * 255 for i in mean], dtype=image.dtype, device=image.device)
        self.std = torch.as_tensor([j * 255 for j in std], dtype=image.dtype, device=image.device)

    def _call_pytorch(self, image):
        # Just in case, make the image a float tensor and apply variance correction.
        image = image.float()
        image.sub_(self.mean[None, :, None, None]).div_(self.std[None, :, None, None])

        return self.model(image)
