0_image_generator_and_emotion_detector.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Required 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 glob\n",
    "from tqdm import tqdm\n",
    "\n",
    "from datasets import augmentations\n",
    "from utils.common import tensor2im, log_input_image\n",
    "from models.psp import pSp\n",
    "\n",
    "from models.stylegan2.model import Generator\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the StyleGANv2 decoder and load pretrained weights\n",
    "\n",
    "style2gan_decoder = Generator(1024, 512, 8) # 1024x1024 output\n",
    "ckpt = torch.load('pretrained_models/stylegan2-ffhq-config-f.pt')\n",
    "style2gan_decoder.load_state_dict(ckpt['g_ema'], strict=False)\n",
    "\n",
    "device = 'cuda'\n",
    "style2gan_decoder.to(device)\n",
    "style2gan_decoder.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Stamina\\.hsemotion\\enet_b0_8_best_afew.pt Compose(\n",
      "    Resize(size=(224, 224), interpolation=bilinear, max_size=None, antialias=None)\n",
      "    ToTensor()\n",
      "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Define the MTCNN face detector for cropping\n",
    "# https://github.com/timesler/facenet-pytorch\n",
    "# https://www.semanticscholar.org/paper/Joint-Face-Detection-and-Alignment-Using-Multitask-Zhang-Zhang/9e60942aa15670ed9ee03af3c0ae011fa4966b7c\n",
    "\n",
    "from facenet_pytorch import MTCNN\n",
    "mtcnn = MTCNN(keep_all=False, select_largest=False, post_process=False, min_face_size=50, device=device)\n",
    "\n",
    "# Define the emotion recognition network\n",
    "# https://github.com/av-savchenko/face-emotion-recognition\n",
    "# https://www.semanticscholar.org/paper/Classifying-Emotions-and-Engagement-in-Online-Based-Savchenko-Savchenko/260d7f95ab8a562f4ff590684ef6a509b8fed316\n",
    "\n",
    "from hsemotion.facial_emotions import HSEmotionRecognizer\n",
    "model_name = 'enet_b0_8_best_afew'\n",
    "fer = HSEmotionRecognizer(model_name=model_name, device=device)\n",
    "emotion_class_n = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function to generate images for given batch \n",
    "\n",
    "# 224x224 image is enough for emotion recognition net\n",
    "# saving 1024x1024 is too slow because of floating point arithmetic during the normalization, so downscale it\n",
    "pool_logit = torch.nn.AdaptiveAvgPool2d((224, 224)) # this is clever, resizes logits (you can not use image resizers, as it is not normalized)\n",
    "\n",
    "def generate_images(generator, n_images):\n",
    "    gen_imgs = []\n",
    "    latents = []\n",
    "    with torch.no_grad():\n",
    "        for _ in range(n_images): # args.pics\n",
    "            z = torch.randn(1, 512, device=device) # args.sample, args.latent\n",
    "            gen_img, latent = generator(\n",
    "                [z], truncation=1, truncation_latent=None, return_latents=True # args.truncation\n",
    "            )\n",
    "            reshaped_gen_img = pool_logit(gen_img)\n",
    "            gen_imgs.append(reshaped_gen_img[0])\n",
    "            latents.append(latent[0].to('cpu').numpy())\n",
    "    return gen_imgs, latents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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_images_and_convert_to_numpy(images): # inplace\n",
    "    for i, image in enumerate(images):\n",
    "        norm_range(image, (-1, 1))\n",
    "        images[i] = image.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to(\"cpu\", torch.uint8).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function to preprocess images and predict emotions for given batch\n",
    "\n",
    "def extract_emotion_scores_from_images(mtcnn_net, images):\n",
    "\n",
    "    # Finding facial bounding box\n",
    "    detections, _ = mtcnn_net.detect(images, landmarks=False)\n",
    "\n",
    "    detections = np.clip(np.array([detection[0] if detection is not None else [-1, -1, -1, -1] for detection in detections]), 0, 223)\n",
    "\n",
    "    # all_scores = []\n",
    "    \n",
    "    # for i, image in enumerate(images):\n",
    "    #     bounding_box = detections[i]\n",
    "    #     box = bounding_box.astype(int)\n",
    "    #     x1, y1, x2, y2 = box[0:4]\n",
    "    #     if x2 == 0:\n",
    "    #         all_scores.append(np.array([-1.] * emotion_class_n, dtype=np.float32)) # TODO Handle no face case !!\n",
    "    #         continue\n",
    "    #     face_img = image[y1:y2,x1:x2,:]\n",
    "    #     _, scores = fer.predict_emotions(face_img, logits=False)\n",
    "    #     all_scores.append(scores)\n",
    "\n",
    "    invalid_images = []\n",
    "\n",
    "    face_images = []\n",
    "    \n",
    "    for i, image in enumerate(images):\n",
    "        bounding_box = detections[i]\n",
    "        box = bounding_box.astype(int)\n",
    "        x1, y1, x2, y2 = box[0:4]\n",
    "        if x2 == 0:\n",
    "            face_images.append(np.zeros((224, 224, 3), np.uint8))\n",
    "            invalid_images.append(i)\n",
    "            continue\n",
    "        face_img = image[y1:y2,x1:x2,:]\n",
    "        face_images.append(face_img)\n",
    "\n",
    "    _, all_scores = fer.predict_multi_emotions(face_images, logits=False)\n",
    "\n",
    "    for invalid_image_idx in invalid_images:\n",
    "        all_scores[invalid_image_idx] = np.zeros_like(all_scores[invalid_image_idx]) - 1\n",
    "    \n",
    "    return all_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function for generates images and extracts predictions for emotions for given batch\n",
    "\n",
    "# {0: 'Anger', 1: 'Contempt', 2: 'Disgust', 3: 'Fear', 4: 'Happiness', 5: 'Neutral', 6: 'Sadness', 7: 'Surprise'}\n",
    "\n",
    "def generate_images_extract_emotions(generator, mtcnn_net, n_images, start_idx = 0, img_save_path = 'main_generatedimages', score_save_path = 'main_emotionscores', latent_save_path = 'main_wplusses'):\n",
    "    generated_images, latents = generate_images(generator, n_images) # max that my gpu can take :(\n",
    "    latents = np.array(latents)\n",
    "    normalize_images_and_convert_to_numpy(generated_images)\n",
    "    scores = extract_emotion_scores_from_images(mtcnn_net, generated_images)\n",
    "\n",
    "    # Save generated images\n",
    "    for i, img in enumerate(generated_images):\n",
    "        img_path = os.path.join(img_save_path, f'{start_idx + i}.png')\n",
    "        img = Image.fromarray(img)\n",
    "        img.save(img_path)\n",
    "    \n",
    "    # Save scores and latents\n",
    "    np.save(os.path.join(score_save_path, f'{start_idx}.npy'), scores)\n",
    "    np.save(os.path.join(latent_save_path, f'{start_idx}.npy'), latents)\n",
    "\n",
    "    # return scores, latents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9776/9776 [8:48:44<00:00,  3.25s/it]  \n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "imgs_to_generate = 312832\n",
    "iterations = imgs_to_generate // batch_size\n",
    "\n",
    "# This is for continuing previous runs\n",
    "start_num = 7168\n",
    "start_batch = start_num // batch_size\n",
    "\n",
    "for i in tqdm(range(iterations)):\n",
    "    generate_images_extract_emotions(style2gan_decoder, mtcnn, batch_size, start_idx = (start_batch + i) * batch_size)"
   ]
  }
 ],
 "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
}