5_hyper_2k_transfer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General imports\n",
    "\n",
    "from argparse import Namespace\n",
    "import time\n",
    "import sys\n",
    "import pprint\n",
    "import numpy as np\n",
    "import os\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision import utils\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import random\n",
    "import glob\n",
    "from tqdm import tqdm\n",
    "import PIL\n",
    "import PIL.Image\n",
    "%matplotlib inline\n",
    "\n",
    "# Importing hyperstyle related\n",
    "\n",
    "hyperstyle_root = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath('__file__')), 'hyperstyle'))\n",
    "sys.path.insert(0, hyperstyle_root)\n",
    "\n",
    "from hyperstyle.utils.model_utils import load_model\n",
    "from hyperstyle.notebooks.notebook_utils import run_alignment\n",
    "from hyperstyle.utils.inference_utils import run_inversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Stamina\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "c:\\Users\\Stamina\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing `weights=ResNet34_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet34_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading HyperStyle from checkpoint: ./pretrained_models/hyperstyle_ffhq.pt\n",
      "Loading pretrained W encoder...\n",
      "Using WEncoder\n",
      "Loading WEncoder from checkpoint: pretrained_models/faces_w_encoder.pt\n"
     ]
    }
   ],
   "source": [
    "# Define the pSp encoder - StyleGANv2 decoder and load pretrained weights\n",
    "\n",
    "device = 'cuda'\n",
    "\n",
    "HYPERSTYLE_ARGS = {\n",
    "    \"hyperstyle_model_path\": \"./pretrained_models/hyperstyle_ffhq.pt\",\n",
    "    \"hyperstyle_w_encoder_path\": \"./pretrained_models/faces_w_encoder.pt\",\n",
    "    \"transform\": transforms.Compose([\n",
    "        transforms.Resize((256, 256)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),\n",
    "}\n",
    "\n",
    "# hyper_net, hyper_opts = load_model(HYPERSTYLE_ARGS['hyperstyle_model_path'], update_opts={'w_encoder_checkpoint_path': HYPERSTYLE_ARGS['hyperstyle_w_encoder_path'], 'resize_outputs': False})\n",
    "hyper_net, hyper_opts = load_model(HYPERSTYLE_ARGS['hyperstyle_model_path'], update_opts={'resize_outputs': False})\n",
    "hyper_net.eval()\n",
    "hyper_net.cuda();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function to normalize images and convert to numpy\n",
    "\n",
    "# This normalization block is taken from the original torch repository:\n",
    "# https://github.com/pytorch/vision/blob/89d2b38cbc3254ed7ed7b43393e4635979ac12eb/torchvision/utils.py\n",
    "\n",
    "def norm_ip(img, low, high):\n",
    "    img.clamp_(min=low, max=high)\n",
    "    img.sub_(low).div_(max(high - low, 1e-5))\n",
    "\n",
    "def norm_range(t, value_range):\n",
    "    if value_range is not None:\n",
    "        norm_ip(t, value_range[0], value_range[1])\n",
    "    else:\n",
    "        norm_ip(t, float(t.min()), float(t.max()))\n",
    "\n",
    "def normalize_image_and_convert_to_numpy(image):\n",
    "    norm_range(image, (-1, 1))\n",
    "    return image.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to(\"cpu\", torch.uint8).numpy()\n",
    "\n",
    "pool_logit = torch.nn.AdaptiveAvgPool2d((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function that generates a random image\n",
    "\n",
    "def generate_random_image(downscale=False):\n",
    "    z = torch.randn(1, 512, device=device) # args.sample, args.latent\n",
    "    gen_img, latent = hyper_net.decoder(\n",
    "        [z], truncation=1, truncation_latent=None, return_latents=True # args.truncation\n",
    "    )\n",
    "    if downscale:\n",
    "        gen_img = pool_logit(gen_img)\n",
    "    return Image.fromarray(normalize_image_and_convert_to_numpy(gen_img[0])), latent # returns tensor\n",
    "\n",
    "# Define a function that generates an image for the given code\n",
    "\n",
    "def generate_image_given_code(code, weights_deltas=None, downscale=False):\n",
    "    with torch.no_grad():\n",
    "        gen_img, latent = hyper_net.decoder(\n",
    "            [code], truncation=1, truncation_latent=None, return_latents=True, input_is_latent=True, randomize_noise=False, weights_deltas=weights_deltas # args.truncation\n",
    "        )\n",
    "        if downscale:\n",
    "            gen_img = pool_logit(gen_img)\n",
    "        return Image.fromarray(normalize_image_and_convert_to_numpy(gen_img[0])), latent, weights_deltas # returns tensor\n",
    "\n",
    "# Define a function that encodes a given image and returns the code alongside its decoding (its 'fake' recreation)\n",
    "\n",
    "def encode_given_image_return_code_and_recreation(image, downscale=False):\n",
    "    with torch.no_grad():\n",
    "        img_transforms = HYPERSTYLE_ARGS['transform']\n",
    "        transformed_image = img_transforms(image)\n",
    "        image, latent, weights_deltas, _ = run_inversion(transformed_image.unsqueeze(0).to(device), hyper_net, hyper_opts, return_intermediate_results=False)\n",
    "        if downscale:\n",
    "            gen_img = pool_logit(gen_img)\n",
    "        return image, latent, weights_deltas\n",
    "\n",
    "# Define a function that generates an image or encodes a given image and returns image-code pair\n",
    "\n",
    "def generate_an_image_code_pair(image=None, downscale=False):\n",
    "    if image is None:\n",
    "        return generate_random_image(downscale)\n",
    "    else:\n",
    "        return encode_given_image_return_code_and_recreation(image, downscale)\n",
    "    \n",
    "# Define a function to load an image from its path, optionally aligns it and returns loadedimage-generatedimage-code\n",
    "\n",
    "def load_image_and_encode(path, predictor=None):\n",
    "    if predictor is not None:\n",
    "        image = run_alignment(path, predictor)\n",
    "    else:\n",
    "        image = PIL.Image.open(path)\n",
    "    gen_image, gen_code, gen_weights_deltas = encode_given_image_return_code_and_recreation(image)\n",
    "    gen_image = Image.fromarray(normalize_image_and_convert_to_numpy(gen_image[0]))\n",
    "    return image, gen_image, gen_code, gen_weights_deltas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to neutralize latent code\n",
    "\n",
    "def neutralize_latent_code(code, neutral_dir, neutral_strength=20):\n",
    "    code, neutral_dir = code.detach().cpu().flatten(), neutral_dir.detach().cpu().flatten()\n",
    "    distance = np.dot(neutral_dir, code) / np.linalg.norm(neutral_dir)\n",
    "    direction = neutral_dir / np.linalg.norm(neutral_dir)\n",
    "    neutral_code = code - distance * direction\n",
    "    neutral_code = neutral_code + neutral_strength * direction\n",
    "    return neutral_code.reshape(18, 512).unsqueeze(0).cuda()\n",
    "\n",
    "# Define a function to transfer emotion from code A to B\n",
    "\n",
    "def transfer_emotion_on_code(code_A, code_B, neutral_dir, neutral_strength):\n",
    "    code_A_neu = neutralize_latent_code(code_A, neutral_dir, neutral_strength)\n",
    "    code_B_neu = neutralize_latent_code(code_B, neutral_dir, neutral_strength)\n",
    "\n",
    "    return (code_A - code_A_neu) + code_B_neu\n",
    "\n",
    "# Define a function to transfer emotion from image A to B\n",
    "\n",
    "def transfer_emotion_on_image(image_A, image_B, neutral_dir, neutral_strength):\n",
    "    _, code_A, wd_A = encode_given_image_return_code_and_recreation(image_A)\n",
    "    _, code_B, wd_B = encode_given_image_return_code_and_recreation(image_B)\n",
    "    \n",
    "    image_A_neu, code_A_neu, wd_A_neu = generate_image_given_code(neutralize_latent_code(code_A, neutral_dir, neutral_strength), wd_A)\n",
    "    image_B_neu, code_B_neu, wd_B_neu = generate_image_given_code(neutralize_latent_code(code_B, neutral_dir, neutral_strength), wd_B)\n",
    "    \n",
    "    image_A_neu_inv, code_A_neu_inv, wd_A_neu_inv = encode_given_image_return_code_and_recreation(image_A_neu) \n",
    "    image_B_neu_inv, code_B_neu_inv, wd_B_neu_inv = encode_given_image_return_code_and_recreation(image_B_neu)\n",
    "\n",
    "    transfer_code = (code_A - code_A_neu_inv) + code_B_neu_inv\n",
    "\n",
    "    return generate_image_given_code(transfer_code, wd_B_neu)[0]\n",
    "\n",
    "# Define a function to transfer emotion from image A to B that utilizes pre-computed delta\n",
    "\n",
    "def transfer_emotion_on_image_using_delta(delta, code_B_neu_inv, wd_B_neu):\n",
    "    return generate_image_given_code(delta + code_B_neu_inv, wd_B_neu)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up directions\n",
    "\n",
    "anger_dir = torch.from_numpy(np.load('main_directions/0.npy').astype(np.float32)).to(device)\n",
    "contempt_dir = torch.from_numpy(np.load('main_directions/1.npy').astype(np.float32)).to(device)\n",
    "disgust_dir = torch.from_numpy(np.load('main_directions/2.npy').astype(np.float32)).to(device)\n",
    "fear_dir = torch.from_numpy(np.load('main_directions/3.npy').astype(np.float32)).to(device)\n",
    "happiness_dir = torch.from_numpy(np.load('main_directions/4.npy').astype(np.float32)).to(device)\n",
    "neutral_dir = torch.from_numpy(np.load('main_directions/5.npy').astype(np.float32)).to(device)\n",
    "sadness_dir = torch.from_numpy(np.load('main_directions/6.npy').astype(np.float32)).to(device)\n",
    "surprise_dir = torch.from_numpy(np.load('main_directions/7.npy').astype(np.float32)).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up basis images and precompute deltas, then save them\n",
    "\n",
    "import pickle\n",
    "emotion_deltas_path = 'emotion_deltas_hyper.pkl'\n",
    "emotions = ['anger', 'contempt', 'disgust', 'fear', 'happiness', 'neutral', 'sadness', 'surprise']\n",
    "\n",
    "# emotion_deltas = dict()\n",
    "\n",
    "# for emotion in emotions:\n",
    "#     _, _, gen_code, gen_wd = load_image_and_encode(f'quant_experiments/basis_aligned/{emotion}_src.jpg')\n",
    "#     image_neu, _, _ = generate_image_given_code(neutralize_latent_code(gen_code, neutral_dir, 20), gen_wd)\n",
    "#     _, gen_code_neu_inv, _ = encode_given_image_return_code_and_recreation(image_neu)\n",
    "#     emotion_deltas[emotion] = (gen_code - gen_code_neu_inv)\n",
    "\n",
    "# with open(emotion_deltas_path, 'wb') as file:\n",
    "#     pickle.dump(emotion_deltas, file)\n",
    "\n",
    "# Load precomputed deltas\n",
    "\n",
    "with open(emotion_deltas_path, 'rb') as file:\n",
    "    emotion_deltas = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dir = 'datasets/celeba_hq_aligned'\n",
    "output_dir_base = 'quant_experiments/transfer_hyper'\n",
    "\n",
    "emotions.append('default')\n",
    "\n",
    "for root, _, files in os.walk(input_dir):\n",
    "    for file in files:\n",
    "        if file.lower().endswith('.jpg'):\n",
    "            file_path = os.path.join(root, file)\n",
    "            _, _, code, wd = load_image_and_encode(file_path)\n",
    "            \n",
    "            image_neu, _, wd_neu = generate_image_given_code(neutralize_latent_code(code, neutral_dir, 20), wd)\n",
    "            _, code_neu_inv, wd_neu_inv = encode_given_image_return_code_and_recreation(image_neu)\n",
    "            for emotion in emotions:\n",
    "                output_dir = os.path.join(output_dir_base, emotion)\n",
    "                relative_path = os.path.relpath(root, input_dir)\n",
    "                output_file_dir = os.path.join(output_dir, relative_path)\n",
    "                if not os.path.exists(output_file_dir):\n",
    "                    os.makedirs(output_file_dir)\n",
    "                output_file_path = os.path.join(output_file_dir, file)\n",
    "\n",
    "                if emotion == 'default':\n",
    "                    image_transfer = generate_image_given_code(code, wd)[0]\n",
    "                else:\n",
    "                    image_transfer = transfer_emotion_on_image_using_delta(emotion_deltas[emotion], code_neu_inv, wd_neu_inv)\n",
    "\n",
    "                image_transfer.save(output_file_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}