import torch
import torchvision
from torchvision import transforms
import requests
from io import BytesIO
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import matplotlib as mpl
import pandas as pd
from torchinfo import summary

model = torchvision.models.mobilenet_v3_small(pretrained=True)

image_url = "https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg"

r = requests.get(image_url)

original_image = Image.open(BytesIO(r.content))

original_image

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

image_tensor = preprocess(original_image).unsqueeze(0).to("cuda")

summary(model.features, input_data=image_tensor)

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
├─ConvBNActivation: 1-1                  [1, 16, 112, 112]         --
|    └─Conv2d: 2-1                       [1, 16, 112, 112]         432
|    └─BatchNorm2d: 2-2                  [1, 16, 112, 112]         32
|    └─Hardswish: 2-3                    [1, 16, 112, 112]         --
├─InvertedResidual: 1-2                  [1, 16, 56, 56]           --
|    └─Sequential: 2-4                   [1, 16, 56, 56]           --
|    |    └─ConvBNActivation: 3-1        [1, 16, 56, 56]           176
|    |    └─SqueezeExcitation: 3-2       [1, 16, 56, 56]           280
|    |    └─ConvBNActivation: 3-3        [1, 16, 56, 56]           288
├─InvertedResidual: 1-3                  [1, 24, 28, 28]           --
|    └─Sequential: 2-5                   [1, 24, 28, 28]           --
|    |    └─ConvBNActivation: 3-4        [1, 72, 56, 56]           1,296
|    |    └─ConvBNActivation: 3-5        [1, 72, 28, 28]           792
|    |    └─ConvBNActivation: 3-6        [1, 24, 28, 28]           1,776
├─InvertedResidual: 1-4                  [1, 24, 28, 28]           --
|    └─Sequential: 2-6                   [1, 24, 28, 28]           --
|    |    └─ConvBNActivation: 3-7        [1, 88, 28, 28]           2,288
|    |    └─ConvBNActivation: 3-8        [1, 88, 28, 28]           968
|    |    └─ConvBNActivation: 3-9        [1, 24, 28, 28]           2,160
├─InvertedResidual: 1-5                  [1, 40, 14, 14]           --
|    └─Sequential: 2-7                   [1, 40, 14, 14]           --
|    |    └─ConvBNActivation: 3-10       [1, 96, 28, 28]           2,496
|    |    └─ConvBNActivation: 3-11       [1, 96, 14, 14]           2,592
|    |    └─SqueezeExcitation: 3-12      [1, 96, 14, 14]           4,728
|    |    └─ConvBNActivation: 3-13       [1, 40, 14, 14]           3,920
├─InvertedResidual: 1-6                  [1, 40, 14, 14]           --
|    └─Sequential: 2-8                   [1, 40, 14, 14]           --
|    |    └─ConvBNActivation: 3-14       [1, 240, 14, 14]          10,080
|    |    └─ConvBNActivation: 3-15       [1, 240, 14, 14]          6,480
|    |    └─SqueezeExcitation: 3-16      [1, 240, 14, 14]          31,024
|    |    └─ConvBNActivation: 3-17       [1, 40, 14, 14]           9,680
├─InvertedResidual: 1-7                  [1, 40, 14, 14]           --
|    └─Sequential: 2-9                   [1, 40, 14, 14]           --
|    |    └─ConvBNActivation: 3-18       [1, 240, 14, 14]          10,080
|    |    └─ConvBNActivation: 3-19       [1, 240, 14, 14]          6,480
|    |    └─SqueezeExcitation: 3-20      [1, 240, 14, 14]          31,024
|    |    └─ConvBNActivation: 3-21       [1, 40, 14, 14]           9,680
├─InvertedResidual: 1-8                  [1, 48, 14, 14]           --
|    └─Sequential: 2-10                  [1, 48, 14, 14]           --
|    |    └─ConvBNActivation: 3-22       [1, 120, 14, 14]          5,040
|    |    └─ConvBNActivation: 3-23       [1, 120, 14, 14]          3,240
|    |    └─SqueezeExcitation: 3-24      [1, 120, 14, 14]          7,832
|    |    └─ConvBNActivation: 3-25       [1, 48, 14, 14]           5,856
├─InvertedResidual: 1-9                  [1, 48, 14, 14]           --
|    └─Sequential: 2-11                  [1, 48, 14, 14]           --
|    |    └─ConvBNActivation: 3-26       [1, 144, 14, 14]          7,200
|    |    └─ConvBNActivation: 3-27       [1, 144, 14, 14]          3,888
|    |    └─SqueezeExcitation: 3-28      [1, 144, 14, 14]          11,704
|    |    └─ConvBNActivation: 3-29       [1, 48, 14, 14]           7,008
├─InvertedResidual: 1-10                 [1, 96, 7, 7]             --
|    └─Sequential: 2-12                  [1, 96, 7, 7]             --
|    |    └─ConvBNActivation: 3-30       [1, 288, 14, 14]          14,400
|    |    └─ConvBNActivation: 3-31       [1, 288, 7, 7]            7,776
|    |    └─SqueezeExcitation: 3-32      [1, 288, 7, 7]            41,832
|    |    └─ConvBNActivation: 3-33       [1, 96, 7, 7]             27,840
├─InvertedResidual: 1-11                 [1, 96, 7, 7]             --
|    └─Sequential: 2-13                  [1, 96, 7, 7]             --
|    |    └─ConvBNActivation: 3-34       [1, 576, 7, 7]            56,448
|    |    └─ConvBNActivation: 3-35       [1, 576, 7, 7]            15,552
|    |    └─SqueezeExcitation: 3-36      [1, 576, 7, 7]            166,608
|    |    └─ConvBNActivation: 3-37       [1, 96, 7, 7]             55,488
├─InvertedResidual: 1-12                 [1, 96, 7, 7]             --
|    └─Sequential: 2-14                  [1, 96, 7, 7]             --
|    |    └─ConvBNActivation: 3-38       [1, 576, 7, 7]            56,448
|    |    └─ConvBNActivation: 3-39       [1, 576, 7, 7]            15,552
|    |    └─SqueezeExcitation: 3-40      [1, 576, 7, 7]            166,608
|    |    └─ConvBNActivation: 3-41       [1, 96, 7, 7]             55,488
├─ConvBNActivation: 1-13                 [1, 576, 7, 7]            --
|    └─Conv2d: 2-15                      [1, 576, 7, 7]            55,296
|    └─BatchNorm2d: 2-16                 [1, 576, 7, 7]            1,152
|    └─Hardswish: 2-17                   [1, 576, 7, 7]            --
==========================================================================================
Total params: 927,008
Trainable params: 927,008
Non-trainable params: 0
Total mult-adds (M): 57.54
==========================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 22.63
Params size (MB): 3.71
Estimated Total Size (MB): 26.94
==========================================================================================

to_pil_image = transforms.ToPILImage()

def get_hidden_layer_output(m, x):
    outputs = []
    with torch.no_grad():
        for layer in m:
            x = layer(x)
            outputs.append(x)
    return outputs

def min_div(n, x):
    x = int(x)
    while True:
        if n % x == 0:
            return  x
        if x > n:
            return n
        x += 1
        
def apply_cmap(image, cmap="viridis"):
    cm = plt.get_cmap(cmap)
    return Image.fromarray(np.uint8(cm(np.array(image.convert("L"))) * 255))

features = get_hidden_layer_output(model.features, image_tensor)

images = []

for i in range(n):
    s = features[i].shape[1]
    nrow = min_div(s, np.sqrt(2 * s))
    images.append(
        apply_cmap(to_pil_image(
        torchvision.utils.make_grid(
    features[i].transpose(0, 1),
            nrow=nrow,
        ))) 
    )

plt.rcParams["font.size"] = 18
fig, axes = plt.subplots(n, figsize=(30, 100))
for i in range(n):
    ax = axes[i]
    image = images[i]
    ax.imshow(image)
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    shape = tuple(features[i].shape)[1:]
    name = model.features[i].__class__.__name__
    title = f"{i+1}: {name}, {shape}"
    ax.set_title(title)