{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d50af37",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import pathlib\n",
    "import matplotlib.pyplot as pyplot\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transformers\n",
    "from models import VGGEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca8cc62f",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "ROOT =Path.cwd()\n",
    "EXAMPLES =ROOT/\"examples\"\n",
    "VGG_WEIGHTS=ROOT/\"models\"/vgg_normalised.ptn\n",
    "device = torch.device(\"cuda\" if torch.cuda.is available() else \"cpu\")\n",
    "image_size = 512\n",
    "\n",
    "image_paths = {\n",
    "    \"Original Brad Pitt\": EXAMPLES / \"brad pitt.jpg\",\n",
    "    \"Stylized Brad Pitt A\": EXAMPLES / \"stylized_brad_pitt.jpg\", \n",
    "    \"Stylized Brad Pitt B\": EXAMPLES / \"stylized_brad pitt (1).jpg\"\n",
    "}\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize( (image_size, image_size)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "encoder = VGGEncoder(str(VGG_WEIGHTS), test=False).to(device).eval()\n",
    "\n",
    "print(\"Using device:\", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70516ffa",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "def load_image(path):\n",
    "    image = Image.open(path).convert(\"RGB\")\n",
    "    tensor = transform(image).unsqueeze(0).to(device)\n",
    "    return image, tensor\n",
    "\n",
    "\n",
    "def show_images (paths_dict):\n",
    "    fig, axes = plt.subplots(1, len (paths_dict), figsize=(16, 5))\n",
    "    for ax, (title, path) in zip (axes, paths_dict.items()):\n",
    "        image = Image.open(path).convert(\"RGB\")\n",
    "        ax.imshow(image)\n",
    "        ax.set_title(title)\n",
    "        ax.axis(\"off\")\n",
    "    plt.tight_layout()\n",
    "\n",
    "\n",
    "def extract_features (tensor):\n",
    "    with torch.no_grad():\n",
    "        h1, h2, h3, h4 encoder (tensor)\n",
    "        return {\n",
    "            \"relul 1\": h1,\n",
    "            \"relu2 1\": h2,\n",
    "            \"relu3 1\": h3,\n",
    "            \"relu4 1\": h4,\n",
    "        }\n",
    "\n",
    "\n",
    "def activation_map (feature_tensor):\n",
    "    activation = feature_tensor[0].mean(dim=0).detach().cpu()\n",
    "    activation = activation - activation.min()\n",
    "    activation = activation / (activation.max() + 1e-8)\n",
    "    return activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd571425",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "image_tensors = {} \n",
    "feature_bank = {}\n",
    "\n",
    "for name, path in image_paths.items():\n",
    "    _, tensor = load_image(path)\n",
    "    image_tensors[name] = tensor\n",
    "    feature_bank[name] = extract_features(tensor)\n",
    "\n",
    "for name, features in feature_bank.items():\n",
    "    shapes = {layer: tuple(feature.shape) for layer, feature in features.items()}\n",
    "    print (name)\n",
    "    print(shapes)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82e8746e",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "layers_to_show = [\"relul_1\", \"relu2_1\", \"relu3_1\", \"relu4_1\"]\n",
    "row_labels = [\"Input Image\"] + layers_to_show\n",
    "num_rows = len(row_labels)\n",
    "num_cols = len(image_paths)\n",
    "\n",
    "fig, axes = plt.subplots (num_rows, num_cols, figsize=(5* num_cols, 3.6 * num_rows))\n",
    "\n",
    "for col, (name, path) in enumerate(image_paths.items()):\n",
    "    image = Image.open (path).convert(\"RGB\")\n",
    "    axes[0, col].imshow(image)\n",
    "    axes[0, col].set_title(name, fontsize=13)\n",
    "    axes [0, col].axis(\"off\")\n",
    "\n",
    "for row, layer in enumerate (layers_to_show, start=1):\n",
    "    axes[row, col].imshow(activation_map(feature_bank[name][layer]))\n",
    "    axes[row, col].axis(\"off\")\n",
    "\n",
    "for row, label in enumerate(row_labels):\n",
    "    axes[row, 0].set_ylabel(label, fontsize=12, rotation=90, labelpad=18)\n",
    "    \n",
    "plt.suptitle(\"VGG Content Features Across Original and Stylized Brad Pitt Images\", fontsize=16, y=1.02)\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}