MWR_experements.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8jS9nXeUUC0E",
        "outputId": "0a9fc541-da7c-429a-dd74-7c4114001b3f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LTV6LT0MdOLt",
        "outputId": "6201f16d-aaf5-4108-8b05-e60e3400dd71"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting pytorchcv\n",
            "  Downloading pytorchcv-0.0.67-py2.py3-none-any.whl (532 kB)\n",
            "\u001b[K     |████████████████████████████████| 532 kB 31.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from pytorchcv) (1.21.6)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from pytorchcv) (2.23.0)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->pytorchcv) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->pytorchcv) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->pytorchcv) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->pytorchcv) (2022.9.24)\n",
            "Installing collected packages: pytorchcv\n",
            "Successfully installed pytorchcv-0.0.67\n"
          ]
        }
      ],
      "source": [
        "!pip install pytorchcv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j_EhkWIGUyN7"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import pandas as pd\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchvision import transforms\n",
        "from torchvision.models.feature_extraction import create_feature_extractor\n",
        "from torchvision.transforms import ToTensor\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "from torchvision.models import vgg16\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "import torch.nn as nn\n",
        "import math \n",
        "from scipy.spatial.distance import euclidean\n",
        "import torch.optim as optim\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
        "from torchvision.transforms import ToPILImage\n",
        "from pytorchcv.model_provider import get_model# as ptcv_get_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DtlAAqykVM8P"
      },
      "outputs": [],
      "source": [
        "def alignment_procedure(img, left_eye, right_eye):\n",
        "    #this function aligns given face in img based on left and right eye coordinates\n",
        "    \n",
        "    left_eye_x, left_eye_y = left_eye\n",
        "    right_eye_x, right_eye_y = right_eye\n",
        "    \n",
        "    #-----------------------\n",
        "    #find rotation direction\n",
        "    \n",
        "    if left_eye_y > right_eye_y:\n",
        "        point_3rd = (right_eye_x, left_eye_y)\n",
        "        direction = -1 #rotate same direction to clock\n",
        "    else:\n",
        "        point_3rd = (left_eye_x, right_eye_y)\n",
        "        direction = 1 #rotate inverse direction of clock\n",
        "    \n",
        "    #-----------------------\n",
        "    #find length of triangle edges\n",
        "    \n",
        "    a = euclidean(left_eye, point_3rd)\n",
        "    b = euclidean(right_eye, point_3rd)\n",
        "    c = euclidean(right_eye, left_eye)\n",
        "    \n",
        "    #-----------------------\n",
        "    \n",
        "    #apply cosine rule\n",
        "    \n",
        "    if b != 0 and c != 0: #this multiplication causes division by zero in cos_a calculation\n",
        "    \n",
        "        cos_a = (b*b + c*c - a*a)/(2*b*c)\n",
        "        angle = np.arccos(cos_a) #angle in radian\n",
        "        angle = (angle * 180) / math.pi #radian to degree\n",
        "    \n",
        "        #-----------------------\n",
        "        #rotate base image\n",
        "    \n",
        "        if direction == -1:\n",
        "            angle = 90 - angle\n",
        "    \n",
        "        img = Image.fromarray(img)\n",
        "        img = np.array(img.rotate(direction * angle))\n",
        "    \n",
        "    #-----------------------\n",
        "    \n",
        "    return img #return img anyway\n",
        "\n",
        "\n",
        "def get_references(features, age: int, df: pd.DataFrame, scheme: str = 'geo', r: float = 0.1, path='')->tuple():\n",
        "  min_age = min(df['age'])\n",
        "  max_age = max(df['age'])\n",
        "  if scheme == 'geo':\n",
        "    c = abs(1 - np.exp(2 * r)) * age / (1 + np.exp(2 * r))\n",
        "    y1 = round(age - c)\n",
        "    y2 = round(age + c)\n",
        "    # y1 = random.randint(3, 25)\n",
        "    # y2 = random.randint(26, 80)\n",
        "    # while y2 - y1 <= 3:\n",
        "    #   y1 = random.randint(3, 25)\n",
        "    #   y2 = random.randint(26, 80)\n",
        "  else:\n",
        "    y1 = int(age - r)\n",
        "    y2 = int(age + r)\n",
        "  if y1 < min_age:\n",
        "    y1 = min_age\n",
        "  if y2 > max_age:\n",
        "    y2 = max_age\n",
        "  # print(y1)\n",
        "  # print('-'*10)\n",
        "  # print(y2)\n",
        "  df_y1 = df[df['age'] == y1]\n",
        "  df_y2 = df[df['age'] == y2]\n",
        "  while len(df_y1) == 0:\n",
        "    #print(f'y1={y1}')\n",
        "    y1 += 1\n",
        "    df_y1 = df[df['age'] == y1]\n",
        "  while len(df_y2) == 0:\n",
        "    #print(f'y2={y2}')\n",
        "    y2 -= 1\n",
        "    df_y2 = df[df['age'] == y2]\n",
        "  # if len(df_y1) == 1:\n",
        "  #   output_y1 = features[df_y1.iloc[0].name]\n",
        "  # else:\n",
        "  #   row_y1 = random.randint(0, len(df_y1) - 1)\n",
        "  #   output_y1 = features[df_y1.iloc[row_y1].name]\n",
        "\n",
        "  # if len(df_y2) == 1:\n",
        "  #   output_y2 = features[df_y2.iloc[0].name]\n",
        "  # else:\n",
        "  #   row_y2 = random.randint(0, len(df_y2) - 1)\n",
        "  #   output_y2 = features[df_y2.iloc[row_y2].name]\n",
        "  if len(df_y1) == 1:\n",
        "    row_y1 = 0\n",
        "  else:\n",
        "    row_y1 = random.randint(0, len(df_y1) - 1)\n",
        "  path_y1 = df_y1.iloc[row_y1, 1]\n",
        "  output_y1, _ = prepare_img(path + path_y1, df_y1, row_y1)\n",
        "\n",
        "  if len(df_y2) == 1:\n",
        "    row_y2 = 0\n",
        "  else:\n",
        "    row_y2 = random.randint(0, len(df_y2) - 1)\n",
        "  path_y2 = df_y2.iloc[row_y2, 1]\n",
        "  output_y2, _ = prepare_img(path + path_y2, df_y2, row_y2)\n",
        "\n",
        "  return output_y1, output_y2, np.log(y1), np.log(y2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jDk9UcXtxJck"
      },
      "outputs": [],
      "source": [
        "def prepare_img(path, df, idx):\n",
        "  #print(path)\n",
        "  img = cv2.imread(path)\n",
        "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "  label = int(df.iloc[idx, 2])\n",
        "  # crop face\n",
        "  bbox = df.iloc[idx, 3][1:-1]\n",
        "  bbox = [int(x) for x in bbox.split(',')]\n",
        "\n",
        "  face_cropped = img[bbox[1]: bbox[3], bbox[0]: bbox[2], :]\n",
        "  # get eyes and align\n",
        "  eye = df.iloc[idx, 4][1:-1]\n",
        "  left_eye = [int(x) for x in eye.split(',')]\n",
        "  left_eye = (left_eye[0], left_eye[1])\n",
        "\n",
        "  eye = df.iloc[idx, 5][1:-1]\n",
        "  right_eye = [int(x) for x in eye.split(',')]\n",
        "  right_eye = (right_eye[0], right_eye[1])\n",
        "  \n",
        "  img = alignment_procedure(face_cropped, left_eye, right_eye)\n",
        "\n",
        "  return img, label"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ug2lsW_sU8-5"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8oKwgd6Helyz"
      },
      "outputs": [],
      "source": [
        "# model = vgg16(pretrained=True)\n",
        "# return_nodes = [\n",
        "#       \"features\"\n",
        "#       ]\n",
        "# model = create_feature_extractor(model, return_nodes=return_nodes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AUvPGWyiVGti"
      },
      "outputs": [],
      "source": [
        "class CustomDataset(Dataset):\n",
        "  def __init__(self, annotations_file, path):\n",
        "      \n",
        "    self.img_labels = pd.read_csv(annotations_file)\n",
        "    self.transform = transforms.Compose([ ToPILImage(), transforms.Resize((224, 224)), ToTensor(),\n",
        "                                          transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "                                        ])\n",
        "    self.path = path\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.img_labels)\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    img_path = str(self.img_labels.iloc[idx, 1])\n",
        "    #img = Image.open(self.path + img_path)\n",
        "    img = cv2.imread(self.path + img_path)\n",
        "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "    label = int(self.img_labels.iloc[idx, 2])\n",
        "    # crop face\n",
        "    bbox = self.img_labels.iloc[idx, 3][1:-1]\n",
        "    bbox = [int(x) for x in bbox.split(',')]\n",
        "\n",
        "    face_cropped = img[bbox[1]: bbox[3], bbox[0]: bbox[2], :]\n",
        "    # get eyes and align\n",
        "    eye = self.img_labels.iloc[idx, 4][1:-1]\n",
        "    left_eye = [int(x) for x in eye.split(',')]\n",
        "    left_eye = (left_eye[0], left_eye[1])\n",
        "\n",
        "    eye = self.img_labels.iloc[idx, 5][1:-1]\n",
        "    right_eye = [int(x) for x in eye.split(',')]\n",
        "    right_eye = (right_eye[0], right_eye[1])\n",
        "    \n",
        "    img = alignment_procedure(face_cropped, left_eye, right_eye)\n",
        "    img = self.transform(img)\n",
        "    #img = img.unsqueeze(0)\n",
        "        \n",
        "    return img, label\n",
        "\n",
        "class CustomDataset2(Dataset):\n",
        "  def __init__(self, annotations_file, path, nns_df, features):\n",
        "      \n",
        "    self.img_labels = pd.read_csv(annotations_file)\n",
        "    #self.transform = transforms.Compose([ ToPILImage(), transforms.Resize((224, 224)), ToTensor(),\n",
        "    #                                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "    #                                    ])\n",
        "    self.path = path\n",
        "    self.features = features\n",
        "    self.nns_df = nns_df\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.img_labels)\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    x = torch.tensor(self.features[idx], device=device, dtype=torch.float)\n",
        "\n",
        "    p_age = int(self.nns_df.iloc[idx, 0])\n",
        "    y1, y2, theta_y1, theta_y2 = get_references(self.features, p_age, self.img_labels)\n",
        "    \n",
        "    label = int(self.img_labels.iloc[idx, 2])\n",
        "    mu = (theta_y1 + theta_y2) / 2\n",
        "    r = (theta_y2 - theta_y1) / 2\n",
        "    if np.log(label) <= theta_y1:\n",
        "      p_rank_gt = -1\n",
        "    elif np.log(label) >= theta_y2:\n",
        "      p_rank_gt = 1\n",
        "    else:\n",
        "      p_rank_gt = (np.log(label) - mu) / (r + 1e-10)\n",
        "    y1 = torch.tensor(y1, device=device, dtype=torch.float)\n",
        "    y2 = torch.tensor(y2, device=device, dtype=torch.float)\n",
        "\n",
        "    return y1.resize_(y1.shape[0], 1, 1), x.resize_(x.shape[0], 1, 1), y2.resize_(y2.shape[0], 1, 1), label, p_rank_gt, mu, r\n",
        "\n",
        "\n",
        "class CustomDataset3(Dataset):\n",
        "  def __init__(self, annotations_file, path, nns_df, features):\n",
        "      \n",
        "    self.img_labels = pd.read_csv(annotations_file)\n",
        "    self.transform = transforms.Compose([ ToPILImage(), transforms.Resize((224, 224)), ToTensor(),\n",
        "                                        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "                                       ])\n",
        "    self.path = path\n",
        "    self.features = features\n",
        "    self.nns_df = nns_df\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.img_labels)\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    # prepare x\n",
        "    img_path = str(self.img_labels.iloc[idx, 1])\n",
        "    \n",
        "    x, theta_x = prepare_img(self.path + img_path, self.img_labels, idx)\n",
        "    \n",
        "    x = self.transform(x)\n",
        "\n",
        "    # prepare y1, y2\n",
        "    p_age = int(self.nns_df.iloc[idx, 0])\n",
        "    y1, y2, theta_y1, theta_y2 = get_references(self.features, p_age, self.img_labels, path = self.path)\n",
        "\n",
        "    mu = (theta_y1 + theta_y2) / 2\n",
        "    r = (theta_y2 - theta_y1) / 2\n",
        "\n",
        "    if np.log(theta_x) <= theta_y1:\n",
        "      p_rank_gt = -1\n",
        "    elif np.log(theta_x) >= theta_y2:\n",
        "      p_rank_gt = 1\n",
        "    else:\n",
        "      p_rank_gt = (np.log(theta_x) - mu) / (r + 1e-10)\n",
        "\n",
        "    y1 = self.transform(y1)\n",
        "    y2 = self.transform(y2)\n",
        "\n",
        "    return y1, x, y2, theta_x, p_rank_gt, mu, r"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-k-en3jKICQo"
      },
      "outputs": [],
      "source": [
        "df_nns_train = pd.read_csv('nns_train_train_after_encoder_checkpoint_v2.csv')#pd.read_csv('nns_train_vgg16.csv')\n",
        "df_nns_val = pd.read_csv('nns_val_train_after_encoder_checkpoint_v2.csv')#pd.read_csv('nns_val_vgg16.csv')\n",
        "df_nns_test = pd.read_csv('nns_test_train_after_encoder_checkpoint_v2.csv')#pd.read_csv('nns_test_vgg16.csv')\n",
        "\n",
        "\n",
        "df_nns_test = df_nns_test.drop(labels='Unnamed: 0', axis=1)\n",
        "df_nns_test['nns'] = df_nns_test['theta']\n",
        "df_nns_test = df_nns_test.drop(labels='theta', axis=1)\n",
        "\n",
        "\n",
        "df_nns_train = df_nns_train.drop(labels='Unnamed: 0', axis=1)\n",
        "df_nns_train['nns'] = df_nns_train['theta']\n",
        "df_nns_train = df_nns_train.drop(labels='theta', axis=1)\n",
        "\n",
        "\n",
        "df_nns_val = df_nns_val.drop(labels='Unnamed: 0', axis=1)\n",
        "df_nns_val['nns'] = df_nns_val['theta']\n",
        "df_nns_val = df_nns_val.drop(labels='theta', axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ub9bn6H5ILQW"
      },
      "outputs": [],
      "source": [
        "train_features = np.load('vgg16_train.npy')\n",
        "val_features = np.load('vgg16_val.npy')\n",
        "test_features = np.load('vgg16_test.npy')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZxqfNiVIVbUE"
      },
      "outputs": [],
      "source": [
        "# ds_train = CustomDataset('train.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Train/Train/')\n",
        "# ds_test = CustomDataset('test.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Test/Test/Test/')\n",
        "# ds_val = CustomDataset('val.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Val/Validation/')\n",
        "\n",
        "ds_train = CustomDataset3('train.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Train/Train/', df_nns_train, train_features)\n",
        "ds_test = CustomDataset3('test.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Test/Test/Test/', df_nns_test, test_features)\n",
        "ds_val = CustomDataset3('val.csv', '/content/drive/MyDrive/AgeEstimation/CLAP/Val/Validation/', df_nns_val, val_features)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vrVXwzAvZkay"
      },
      "outputs": [],
      "source": [
        "class Regressor(nn.Sequential):\n",
        "    def __init__(self, input_channel, output_channel):\n",
        "      super(Regressor, self).__init__()\n",
        "      self.convA = nn.Conv2d(input_channel, output_channel, kernel_size=1, stride=1)\n",
        "      self.leakyreluA = nn.ReLU()\n",
        "      self.convB = nn.Conv2d(output_channel, output_channel, kernel_size=1, stride=1)\n",
        "      self.leakyreluB = nn.ReLU()\n",
        "      self.dropout = nn.Dropout(p=0.5)\n",
        "      self.convC = nn.Conv2d(output_channel, 1, kernel_size=1, stride=1)\n",
        "      self.activation = nn.Tanh()\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "      x = self.convA(x)\n",
        "      x = self.leakyreluA(x)\n",
        "      x = self.convB(x)\n",
        "      x = self.leakyreluB(x)\n",
        "      x = self.dropout(x)\n",
        "      x = self.convC(x)\n",
        "\n",
        "      return self.activation(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xA2C2s-pb98J"
      },
      "outputs": [],
      "source": [
        "class Global_Regressor(nn.Module):\n",
        "    def __init__(self):\n",
        "      super(Global_Regressor, self).__init__()\n",
        "      #model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)\n",
        "      # model = vgg16(pretrained=True)\n",
        "      # return_nodes = [\n",
        "      # \"features\"\n",
        "      # ]\n",
        "      # model = create_feature_extractor(model, return_nodes=return_nodes)\n",
        "      self.encoder = get_model(\"bn_vgg16\", pretrained=True)\n",
        "      #self.encoder = model\n",
        "      self.avg_pool = nn.AvgPool2d(kernel_size=7)\n",
        "      self.regressor = Regressor(1536, 512)\n",
        "\n",
        "    def forward_siamese(self, x):\n",
        "      x = self.encoder.features.stage1(x)\n",
        "      x = self.encoder.features.stage2(x)\n",
        "      x = self.encoder.features.stage3(x)\n",
        "      x = self.encoder.features.stage4(x)\n",
        "      x = self.encoder.features.stage5(x)\n",
        "      x = self.avg_pool(x)\n",
        "\n",
        "      return x\n",
        "\n",
        "    def forward(self, y1, x, y2, stage='train'):\n",
        "      \n",
        "      if stage == 'etc':\n",
        "        x = torch.cat([y1, x, y2], dim=1)\n",
        "\n",
        "        output = self.regressor(x)\n",
        "      else:\n",
        "        y1 = self.forward_siamese(y1)\n",
        "        x = self.forward_siamese(x)\n",
        "        y2 = self.forward_siamese(y2)\n",
        "\n",
        "        x = torch.cat([y1, x, y2], dim=1)\n",
        "\n",
        "        output = self.regressor(x)\n",
        "\n",
        "      return output\n",
        "\n",
        "      "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t21fjwkF2eIy"
      },
      "outputs": [],
      "source": [
        "class Global_Regressor2(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(Global_Regressor2, self).__init__()\n",
        "    #model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)\n",
        "    # model = vgg16(pretrained=True)\n",
        "    return_nodes = [\n",
        "    \"features\"\n",
        "    ]\n",
        "    # model = create_feature_extractor(model, return_nodes=return_nodes)\n",
        "    self.encoder = create_feature_extractor(get_model(\"bn_vgg16\", pretrained=True), return_nodes=return_nodes)\n",
        "    #self.encoder = model\n",
        "    self.avg_pool = nn.AvgPool2d(kernel_size=7)\n",
        "    self.regressor = Regressor(1536, 512)\n",
        "\n",
        "  def forward_siamese(self, x):\n",
        "    x = self.encoder(x)['features']\n",
        "    # x = self.encoder.features.stage2(x)\n",
        "    # x = self.encoder.features.stage3(x)\n",
        "    # x = self.encoder.features.stage4(x)\n",
        "    # x = self.encoder.features.stage5(x)\n",
        "    x = self.avg_pool(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "  def forward(self, y1, x, y2, stage='train'):\n",
        "    if stage == 'etc':\n",
        "      x = torch.cat([y1, x, y2], dim=1)\n",
        "\n",
        "      output = self.regressor(x)\n",
        "    else:\n",
        "      y1 = self.forward_siamese(y1)\n",
        "      x = self.forward_siamese(x)\n",
        "      y2 = self.forward_siamese(y2)\n",
        "\n",
        "      x = torch.cat([y1, x, y2], dim=1)\n",
        "\n",
        "      output = self.regressor(x)\n",
        "\n",
        "    return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EVr2Ncgcfix4",
        "outputId": "3dfb807f-6c3a-41f1-e0e6-7985844cd208"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading /root/.torch/models/bn_vgg16-0779-0f570b92.pth.zip from https://github.com/osmr/imgclsmob/releases/download/v0.0.359/bn_vgg16-0779-0f570b92.pth.zip...\n"
          ]
        }
      ],
      "source": [
        "model = Global_Regressor2()\n",
        "model = model.to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8w9boh092jn0",
        "outputId": "e038209b-b3a9-4391-b43f-9cb852dfdda6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.load_state_dict(torch.load('drive/MyDrive/model_MWR_global.pth'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "62k6cmTArd7U"
      },
      "outputs": [],
      "source": [
        "batch_size = 32"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z1s81PRhi94e"
      },
      "outputs": [],
      "source": [
        "train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=False, num_workers=1)\n",
        "val_loader = DataLoader(ds_val, batch_size=batch_size, shuffle=False, num_workers=1)\n",
        "test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False, num_workers=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JrXM_q5FqDud"
      },
      "outputs": [],
      "source": [
        "# model = vgg16(pretrained=True)\n",
        "# return_nodes = [\n",
        "#       \"features\"\n",
        "#       ]\n",
        "# model = create_feature_extractor(model, return_nodes=return_nodes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vizJRz02qF1g"
      },
      "outputs": [],
      "source": [
        "# train_features = []\n",
        "# val_features = []\n",
        "# test_features = []\n",
        "# pool = nn.AvgPool2d(kernel_size=7)\n",
        "\n",
        "\n",
        "# for step,batch in enumerate(train_loader):\n",
        "#   if step % 50 == 0 and not step == 0:\n",
        "#     print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(train_loader)))\n",
        "#   images, labels = batch\n",
        "#   with torch.no_grad():\n",
        "#     features = model(images)['features']\n",
        "#     features = pool(features)\n",
        "#   for i in range(len(features)):\n",
        "#     b = features[i].squeeze().detach().cpu().numpy()\n",
        "#     train_features.append(b)\n",
        "\n",
        "\n",
        "\n",
        "# for step,batch in enumerate(val_loader):\n",
        "#   if step % 10 == 0 and not step == 0:\n",
        "#     print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(val_loader)))\n",
        "#   images, labels = batch\n",
        "#   with torch.no_grad():\n",
        "#     features = model(images)['features']\n",
        "#     features = pool(features)\n",
        "#   for i in range(len(features)):\n",
        "#     b = features[i].squeeze().detach().cpu().numpy()\n",
        "#     val_features.append(b)\n",
        "\n",
        "\n",
        "# for step,batch in enumerate(test_loader):\n",
        "#   if step % 10 == 0 and not step == 0:\n",
        "#     print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(test_loader)))\n",
        "#   images, labels = batch\n",
        "#   with torch.no_grad():\n",
        "#     features = model(images)['features']\n",
        "#     features = pool(features)\n",
        "#   for i in range(len(features)):\n",
        "#     b = features[i].squeeze().detach().cpu().numpy()\n",
        "#     test_features.append(b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EYgi6K7nsJ5B"
      },
      "outputs": [],
      "source": [
        "# train_features = np.array(train_features)\n",
        "# val_features = np.array(val_features)\n",
        "# test_features = np.array(test_features)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m1dV7RoKsT7K"
      },
      "outputs": [],
      "source": [
        "# np.save('vgg16_train.npy', train_features)\n",
        "# np.save('vgg16_val.npy', val_features)\n",
        "# np.save('vgg16_test.npy', test_features)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OzadYyA-6mbb"
      },
      "outputs": [],
      "source": [
        "# df_train = pd.read_csv('train.csv')\n",
        "# df_val = pd.read_csv('val.csv')\n",
        "# df_test = pd.read_csv('test.csv')\n",
        "# df_train.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A82BJ33B5ffM"
      },
      "outputs": [],
      "source": [
        "# from sklearn.metrics import pairwise_distances\n",
        "\n",
        "# train_ages = list(df_train['age'])\n",
        "# val_ages = list(df_val['age'])\n",
        "# test_ages = list(df_test['age'])\n",
        "\n",
        "# train_theta = []\n",
        "# for i, e in enumerate(train_features):\n",
        "#   e = e.reshape(1, -1)\n",
        "#   dists = []\n",
        "#   for j, e2 in enumerate(train_features):\n",
        "#     e2 = e2.reshape(1, -1)\n",
        "#     dists.append(pairwise_distances(e, e2))\n",
        "  \n",
        "#   ages = train_ages.copy()\n",
        "\n",
        "#   dists.pop(i)\n",
        "#   ages.pop(i)\n",
        "  \n",
        "#   dists, ages = zip(*sorted(zip(dists, ages), reverse=True))\n",
        "#   p = sum(ages[:5]) / 5\n",
        "#   train_theta.append(p)\n",
        "\n",
        "# val_theta = []\n",
        "# for i, e in enumerate(val_features):\n",
        "#   e = e.reshape(1, -1)\n",
        "#   dists = []\n",
        "#   for j, e2 in enumerate(val_features):\n",
        "#     e2= e2.reshape(1, -1)\n",
        "#     dists.append(pairwise_distances(e, e2))\n",
        "  \n",
        "#   ages = val_ages.copy()\n",
        "\n",
        "#   dists.pop(i)\n",
        "#   ages.pop(i)\n",
        "  \n",
        "#   dists, ages = zip(*sorted(zip(dists, ages), reverse=True))\n",
        "#   p = sum(ages[:5]) / 5\n",
        "#   val_theta.append(p)\n",
        "\n",
        "\n",
        "# test_theta = []\n",
        "# for i, e in enumerate(test_features):\n",
        "#   e = e.reshape(1, -1)\n",
        "#   dists = []\n",
        "#   for j, e2 in enumerate(test_features):\n",
        "#     e2 = e2.reshape(1, -1)\n",
        "#     dists.append(pairwise_distances(e, e2))\n",
        "  \n",
        "#   ages = test_ages.copy()\n",
        "\n",
        "#   dists.pop(i)\n",
        "#   ages.pop(i)\n",
        "  \n",
        "#   dists, ages = zip(*sorted(zip(dists, ages), reverse=True))\n",
        "#   p = sum(ages[:5]) / 5\n",
        "#   test_theta.append(p)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E9fd565p9Pez"
      },
      "outputs": [],
      "source": [
        "# df_nns_train = pd.DataFrame(data={'nns': train_theta})\n",
        "# df_nns_val = pd.DataFrame(data={'nns': val_theta})\n",
        "# df_nns_test = pd.DataFrame(data={'nns': test_theta})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M8ZjQ8CL91wx"
      },
      "outputs": [],
      "source": [
        "# df_nns_train.to_csv('nns_train_vgg16.csv', index=False)\n",
        "# df_nns_val.to_csv('nns_val_vgg16.csv', index=False)\n",
        "# df_nns_test.to_csv('nns_test_vgg16.csv', index=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "272O2-BVky3e"
      },
      "outputs": [],
      "source": [
        "optimizer = optim.AdamW(model.parameters(), lr=0.0001)\n",
        "scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=2,verbose=True)\n",
        "criterion = nn.MSELoss()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B3vTg60zlrDK"
      },
      "outputs": [],
      "source": [
        "epochs = 300\n",
        "df_test = pd.read_csv('test.csv')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s90GwcWWkjPc"
      },
      "outputs": [],
      "source": [
        "def train():\n",
        "  \n",
        "  model.train()\n",
        "\n",
        "  total_loss = 0.0\n",
        "\n",
        "  total_preds=[]\n",
        "\n",
        "  for step,batch in enumerate(train_loader):\n",
        "    #if step % 50 == 0 and not step == 0:\n",
        "    print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(train_loader)))\n",
        "    batch = [r.to(device) for r in batch]\n",
        "    \n",
        "    y1, x, y2, labels, p_rank_gts, _, _ = batch\n",
        "    p_rank_gts = torch.tensor(p_rank_gts, device=device, dtype=torch.float)\n",
        "    model.zero_grad() \n",
        "    preds = model(y1, x, y2).squeeze()\n",
        "    loss = criterion(preds, p_rank_gts)\n",
        "    total_loss = total_loss + loss.item()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    preds=preds.detach().cpu().numpy()\n",
        "    total_preds.append(preds)\n",
        "\n",
        "  avg_loss = total_loss / len(train_loader)\n",
        "\n",
        "  total_preds  = np.concatenate(total_preds, axis=0)\n",
        "\n",
        "  return avg_loss, total_preds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fzzIkfQhks_F"
      },
      "outputs": [],
      "source": [
        "transform = transforms.Compose([ ToPILImage(), transforms.Resize((224, 224)), ToTensor(),\n",
        "                                        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "                                       ])\n",
        "\n",
        "def evaluate(loader, stage='val'):\n",
        "  \n",
        "  print(\"\\nEvaluating...\")\n",
        "\n",
        "  model.eval()\n",
        "\n",
        "  total_loss = 0.0\n",
        "\n",
        "  total_preds = []\n",
        "  d = {}\n",
        "  for step,batch in enumerate(loader):\n",
        "  \n",
        "    #if step % 50 == 0 and not step == 0:\n",
        "              \n",
        "    print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(loader)))\n",
        "\n",
        "    batch = [t.to(device) for t in batch]\n",
        "    \n",
        "    #faces, fulls, labels = batch\n",
        "    y1, x, y2, labels, p_rank_gts, mu, r = batch\n",
        "    p_rank_gts = torch.tensor(p_rank_gts, device=device, dtype=torch.float)\n",
        "    #r = r.detach().cpu().numpy()\n",
        "    #mu = r.detach().cpu().numpy()\n",
        "    with torch.no_grad():\n",
        "  \n",
        "      #preds = model(faces, fulls)\n",
        "      preds = model(y1, x, y2).squeeze()\n",
        "\n",
        "      loss = criterion(preds,p_rank_gts)\n",
        "\n",
        "      total_loss = total_loss + loss.item()\n",
        "\n",
        "      preds = preds.detach().cpu().numpy()\n",
        "      thetas = []\n",
        "      for p, r_curr, mu_curr in zip(preds, r, mu):\n",
        "        elem = np.exp(p * r_curr.detach().cpu().item() + mu_curr.detach().cpu().item())\n",
        "        thetas.append(elem)\n",
        "      if 0 not in d.keys():\n",
        "        d[0] = []\n",
        "      if stage != 'val':\n",
        "        d[0].append(thetas)\n",
        "        for i in range(4):\n",
        "          y1s = []\n",
        "          y2s = []\n",
        "          mus = []\n",
        "          rs = []\n",
        "          for e in thetas:\n",
        "            \n",
        "            y1, y2, theta_y1, theta_y2 = get_references(test_features, e, df_test, path='/content/drive/MyDrive/AgeEstimation/CLAP/Test/Test/Test/')\n",
        "            \n",
        "            mu_i = (theta_y1 + theta_y2) / 2\n",
        "            r_i = (theta_y2 - theta_y1) / 2\n",
        "            y1s.append(transform(y1))\n",
        "            y2s.append(transform(y2)) \n",
        "            mus.append(mu_i) \n",
        "            rs.append(r_i) \n",
        "          #y1s = np.array(y1s)\n",
        "          #y2s = np.array(y2s)\n",
        "          # y1 = transform(y1s)\n",
        "          # y2 = transform(y2s)\n",
        "          y1 = torch.stack(y1s).to(device)#.resize_(y1s.shape[0], y1s.shape[1], 1, 1)\n",
        "          y2 = torch.stack(y2s).to(device)#.resize_(y2s.shape[0], y2s.shape[1], 1, 1)\n",
        "          preds = model(y1, x, y2).squeeze()\n",
        "          preds = preds.detach().cpu().numpy()\n",
        "\n",
        "          thetas = []\n",
        "          for p, r_curr, mu_curr in zip(preds, rs, mus):\n",
        "            thetas.append(np.exp(p * r_curr + mu_curr))\n",
        "          if i + 1 not in d.keys():\n",
        "            d[i + 1] = []\n",
        "          \n",
        "          d[i + 1].append(thetas)\n",
        "\n",
        "      total_preds.append(thetas)\n",
        "\n",
        "  avg_loss = total_loss / len(loader) \n",
        "\n",
        "  total_preds  = np.concatenate(total_preds, axis=0)\n",
        "\n",
        "  return avg_loss, total_preds, d"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "Bo3ojJQuleqW"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import mean_absolute_error\n",
        "# set initial loss to infinite\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "# empty lists to store training and validation loss of each epoch\n",
        "train_losses=[]\n",
        "valid_losses=[]\n",
        "test_losses=[]\n",
        "\n",
        "#for each epoch\n",
        "for epoch in range(epochs):\n",
        "     \n",
        "    print('\\n Epoch {:} / {:}'.format(epoch + 1, epochs))\n",
        "    train_loss, _ = train()\n",
        "\n",
        "    valid_loss, _, _ = evaluate(val_loader)\n",
        "    scheduler.step(valid_loss)\n",
        "\n",
        "    if valid_loss < best_valid_loss:\n",
        "       best_valid_loss = valid_loss\n",
        "        \n",
        "    #torch.save(model.state_dict(), 'drive/MyDrive/MWR_CLAP' + str(epoch) + '.pt')\n",
        "\n",
        "    train_losses.append(train_loss)\n",
        "    valid_losses.append(valid_loss)\n",
        "    \n",
        "    \n",
        "    # with open('drive/MyDrive/train_loss.txt', 'a') as f:\n",
        "    #     f.write(str(train_loss) + '\\n')\n",
        "        \n",
        "        \n",
        "    # with open('drive/MyDrive/valid_loss.txt', 'a') as f:\n",
        "    #     f.write(str(valid_loss) + '\\n')\n",
        "\n",
        "        \n",
        "    test_loss, preds, list_preds = evaluate(test_loader, stage='test')\n",
        "    test_losses.append(test_loss)\n",
        "    \n",
        "    \n",
        "    # with open('drive/MyDrive/test_loss.txt', 'a') as f:\n",
        "    #     f.write(str(test_loss) + '\\n')\n",
        "        \n",
        "\n",
        "    print(f'\\nTraining Loss: {train_loss:.3f}')\n",
        "    print(f'Validation Loss: {valid_loss:.3f}')\n",
        "\n",
        "    for i in range(len(list_preds.keys())):\n",
        "      \n",
        "      MAE = mean_absolute_error(list(df_test['age']), np.concatenate(list_preds[i], axis=0))\n",
        "      print(f'MAE on MWR with {i + 1} steps = {MAE}')\n",
        "\n",
        "    print(f'\\nMAE on test-set: {MAE:.3f}')\n",
        "\n",
        "print(f'\\nBest valid loss: {best_valid_loss:.3f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CmYQGSA0TRXA"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}