import pandas as pd
import torch
from torch import nn
from torchvision import models, transforms

#support for resnet-18, resnet-34, alexnet. Others yet to be done
#support for resnext and wide resnet added
#removed resizing of images

class Img2Vec():
    def __init__(self, cuda = False, model = "resnet-18", layer = "default", layer_output_size = 512):
        
        self.device = torch.device("cuda" if cuda else "cpu")
        self.layer_output_size = layer_output_size
        self.model_name = model
        
        self.model, self.extraction_layer = self._get_model_and_layer(model, layer)
        
        self.model = self.model.to(self.device)
        self.model.eval()
        
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])
        
    def get_vec(self, img, tensor = False):
        image = self.normalize(self.to_tensor(img)).unsqueeze(0).to(self.device)

        if self.model_name == 'alexnet':
            my_embedding = torch.zeros(1, self.layer_output_size)
        else:
            my_embedding = torch.zeros(1, self.layer_output_size, 1, 1)

        def copy_data(m, i, o):
            my_embedding.copy_(o.data)

        h = self.extraction_layer.register_forward_hook(copy_data)
        h_x = self.model(image)
        h.remove()

        if tensor:
            return my_embedding
        else:
            if self.model_name == 'alexnet':
                return my_embedding.numpy()[0, :]
            else:
                return my_embedding.numpy()[0, :, 0, 0]

    def _get_model_and_layer(self, model_name, layer):
        
        if model_name == 'resnet18':
            model = models.resnet18(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)

            return model, layer

        elif model_name == 'alexnet':
            model = models.alexnet(pretrained=True)
            if layer == 'default':
                layer = model.classifier[-2]
                self.layer_output_size = 4096
            else:
                layer = model.classifier[-layer]

            return model, layer
        
        elif model_name == 'resnet34':
            model = models.resnet34(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)

            return model, layer
        
        elif model_name == "wide_resnet":
            model = models.wide_resnet50_2(pretrained = True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            
            return model, layer
        
        elif model_name == "resnext":
            model = models.resnext50_32x4d(pretrained = True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            
            return model, layer

        else:
            raise KeyError('Model %s was not found' % model_name)
7ed9e92bf82c06f629fae8cba65d28bd962ce67c
<mask token>


class Img2Vec:
    <mask token>
    <mask token>
    <mask token>

<mask token>


class Img2Vec:
    <mask token>

    def get_vec(self, img, tensor=False):
        image = self.normalize(self.to_tensor(img)).unsqueeze(0).to(self.device
            )
        if self.model_name == 'alexnet':
            my_embedding = torch.zeros(1, self.layer_output_size)
        else:
            my_embedding = torch.zeros(1, self.layer_output_size, 1, 1)

        def copy_data(m, i, o):
            my_embedding.copy_(o.data)
        h = self.extraction_layer.register_forward_hook(copy_data)
        h_x = self.model(image)
        h.remove()
        if tensor:
            return my_embedding
        elif self.model_name == 'alexnet':
            return my_embedding.numpy()[0, :]
        else:
            return my_embedding.numpy()[0, :, 0, 0]

    def _get_model_and_layer(self, model_name, layer):
        if model_name == 'resnet18':
            model = models.resnet18(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'alexnet':
            model = models.alexnet(pretrained=True)
            if layer == 'default':
                layer = model.classifier[-2]
                self.layer_output_size = 4096
            else:
                layer = model.classifier[-layer]
            return model, layer
        elif model_name == 'resnet34':
            model = models.resnet34(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'wide_resnet':
            model = models.wide_resnet50_2(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'resnext':
            model = models.resnext50_32x4d(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        else:
            raise KeyError('Model %s was not found' % model_name)

<mask token>


class Img2Vec:

    def __init__(self, cuda=False, model='resnet-18', layer='default',
        layer_output_size=512):
        self.device = torch.device('cuda' if cuda else 'cpu')
        self.layer_output_size = layer_output_size
        self.model_name = model
        self.model, self.extraction_layer = self._get_model_and_layer(model,
            layer)
        self.model = self.model.to(self.device)
        self.model.eval()
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225])

    def get_vec(self, img, tensor=False):
        image = self.normalize(self.to_tensor(img)).unsqueeze(0).to(self.device
            )
        if self.model_name == 'alexnet':
            my_embedding = torch.zeros(1, self.layer_output_size)
        else:
            my_embedding = torch.zeros(1, self.layer_output_size, 1, 1)

        def copy_data(m, i, o):
            my_embedding.copy_(o.data)
        h = self.extraction_layer.register_forward_hook(copy_data)
        h_x = self.model(image)
        h.remove()
        if tensor:
            return my_embedding
        elif self.model_name == 'alexnet':
            return my_embedding.numpy()[0, :]
        else:
            return my_embedding.numpy()[0, :, 0, 0]

    def _get_model_and_layer(self, model_name, layer):
        if model_name == 'resnet18':
            model = models.resnet18(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'alexnet':
            model = models.alexnet(pretrained=True)
            if layer == 'default':
                layer = model.classifier[-2]
                self.layer_output_size = 4096
            else:
                layer = model.classifier[-layer]
            return model, layer
        elif model_name == 'resnet34':
            model = models.resnet34(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'wide_resnet':
            model = models.wide_resnet50_2(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'resnext':
            model = models.resnext50_32x4d(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        else:
            raise KeyError('Model %s was not found' % model_name)

import pandas as pd
import torch
from torch import nn
from torchvision import models, transforms


class Img2Vec:

    def __init__(self, cuda=False, model='resnet-18', layer='default',
        layer_output_size=512):
        self.device = torch.device('cuda' if cuda else 'cpu')
        self.layer_output_size = layer_output_size
        self.model_name = model
        self.model, self.extraction_layer = self._get_model_and_layer(model,
            layer)
        self.model = self.model.to(self.device)
        self.model.eval()
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225])

    def get_vec(self, img, tensor=False):
        image = self.normalize(self.to_tensor(img)).unsqueeze(0).to(self.device
            )
        if self.model_name == 'alexnet':
            my_embedding = torch.zeros(1, self.layer_output_size)
        else:
            my_embedding = torch.zeros(1, self.layer_output_size, 1, 1)

        def copy_data(m, i, o):
            my_embedding.copy_(o.data)
        h = self.extraction_layer.register_forward_hook(copy_data)
        h_x = self.model(image)
        h.remove()
        if tensor:
            return my_embedding
        elif self.model_name == 'alexnet':
            return my_embedding.numpy()[0, :]
        else:
            return my_embedding.numpy()[0, :, 0, 0]

    def _get_model_and_layer(self, model_name, layer):
        if model_name == 'resnet18':
            model = models.resnet18(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'alexnet':
            model = models.alexnet(pretrained=True)
            if layer == 'default':
                layer = model.classifier[-2]
                self.layer_output_size = 4096
            else:
                layer = model.classifier[-layer]
            return model, layer
        elif model_name == 'resnet34':
            model = models.resnet34(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'wide_resnet':
            model = models.wide_resnet50_2(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        elif model_name == 'resnext':
            model = models.resnext50_32x4d(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            return model, layer
        else:
            raise KeyError('Model %s was not found' % model_name)

import pandas as pd
import torch
from torch import nn
from torchvision import models, transforms

#support for resnet-18, resnet-34, alexnet. Others yet to be done
#support for resnext and wide resnet added
#removed resizing of images

class Img2Vec():
    def __init__(self, cuda = False, model = "resnet-18", layer = "default", layer_output_size = 512):
        
        self.device = torch.device("cuda" if cuda else "cpu")
        self.layer_output_size = layer_output_size
        self.model_name = model
        
        self.model, self.extraction_layer = self._get_model_and_layer(model, layer)
        
        self.model = self.model.to(self.device)
        self.model.eval()
        
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])
        
    def get_vec(self, img, tensor = False):
        image = self.normalize(self.to_tensor(img)).unsqueeze(0).to(self.device)

        if self.model_name == 'alexnet':
            my_embedding = torch.zeros(1, self.layer_output_size)
        else:
            my_embedding = torch.zeros(1, self.layer_output_size, 1, 1)

        def copy_data(m, i, o):
            my_embedding.copy_(o.data)

        h = self.extraction_layer.register_forward_hook(copy_data)
        h_x = self.model(image)
        h.remove()

        if tensor:
            return my_embedding
        else:
            if self.model_name == 'alexnet':
                return my_embedding.numpy()[0, :]
            else:
                return my_embedding.numpy()[0, :, 0, 0]

    def _get_model_and_layer(self, model_name, layer):
        
        if model_name == 'resnet18':
            model = models.resnet18(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)

            return model, layer

        elif model_name == 'alexnet':
            model = models.alexnet(pretrained=True)
            if layer == 'default':
                layer = model.classifier[-2]
                self.layer_output_size = 4096
            else:
                layer = model.classifier[-layer]

            return model, layer
        
        elif model_name == 'resnet34':
            model = models.resnet34(pretrained=True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)

            return model, layer
        
        elif model_name == "wide_resnet":
            model = models.wide_resnet50_2(pretrained = True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            
            return model, layer
        
        elif model_name == "resnext":
            model = models.resnext50_32x4d(pretrained = True)
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 2048
            else:
                layer = model._modules.get(layer)
            
            return model, layer

        else:
            raise KeyError('Model %s was not found' % model_name)