RAFDBClass.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Untitled1.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyPVGAeiwG5g9o9kaOUnx+jZ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/SridharSola/Contrastive-learning-FER/blob/main/RAFDBClass.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TZrpcRjj5fwZ"
      },
      "outputs": [],
      "source": [
        "#Imports\n",
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data\n",
        "import torch.optim\n",
        "import os #for creating and removing directories\n",
        "import torch.utils.data as data\n",
        "import matplotlib.pyplot as plt #to view the image as images are defaulted to objects of Pillow lib in Python\n",
        "%matplotlib inline \n",
        "import torchvision.transforms as transforms\n",
        "#import torchvision.transforms.functional as TF #to convert images to tensors\n",
        "import argparse\n",
        "from PIL import Image\n",
        "import pandas as pd\n",
        "import cv2\n",
        "from torch.utils.data import random_split\n",
        "from IPython import display\n",
        "import random\n",
        "import albumentations as A #For applying same transforms\n",
        "from torchvision.datasets.folder import default_loader\n",
        "\n",
        "class RandomChoice(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.t = random.choice(self.transforms)\n",
        "\n",
        "    def __call__(self, img):\n",
        "        return self.t(img)\n",
        "\n",
        "class Contr_Learn_RafDataset(data.Dataset):\n",
        "  def __init__(self, root, mask_file, partition = 'train', transform = None, num_classes = 7, loader = default_loader, UU = False, MM = False):\n",
        "    \"\"\"\n",
        "      root --> path to ocation of images in drive\n",
        "      mask_file --> path for masked RAF-DB images\n",
        "      num_classes --> number of annotations (7 in RAF-DB)\n",
        "      labels are same for bothe masked and unmasked\n",
        "      Need mask label as well for mask detection task\n",
        "\n",
        "      Note: We read the labels here but leave reading of images to __getitem__\n",
        "    \"\"\"\n",
        "    self.root = root\n",
        "    self.mask_file = mask_file\n",
        "    self.num_classes = num_classes\n",
        "    self.partition = partition\n",
        "    self.loader = loader\n",
        "    self.transform = transforms.Compose([transforms.ToPILImage(),\n",
        "                                        transforms.Resize((224, 224)),\n",
        "                                        transforms.ToTensor(),\n",
        "                                        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
        "                                                            std=[0.229, 0.224, 0.225])])\n",
        "\n",
        "    if partition == 'train':\n",
        "      \n",
        "      \n",
        "      self.transform1 = A.Compose([\n",
        "                                 A.HorizontalFlip(), A.HueSaturationValue(), A.RandomContrast(), \n",
        "                                 A.ShiftScaleRotate(shift_limit = 0.0625,scale_limit = 0.1 ,rotate_limit = 3, p = 0.5),\n",
        "                                 A.IAAAffine(scale = (1.0, 1.25), rotate = 0.0, p = 0.5)\n",
        "                                 ],\n",
        "                                  additional_targets={'image1':'image'})\n",
        "      self.transform2 = transforms.Compose([\n",
        "                                            transforms.ToPILImage(),\n",
        "                                            transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n",
        "                                            ])\n",
        "     \n",
        "\n",
        "\n",
        "    NAME_COLUMN = 0\n",
        "    LABEL_COLUMN = 1\n",
        "    #Reading labels\n",
        "    train_txtfile = pd.read_csv(os.path.join(self.root, 'train_label.txt'), sep=' ', header=None)\n",
        "    test_txtfile = pd.read_csv(os.path.join(self.root, 'test_label.txt'), sep=' ', header=None)\n",
        "    if partition == 'train':\n",
        "      self.label = train_txtfile.iloc[:, LABEL_COLUMN].values - 1 # 0:Surprise, 1:Fear, 2:Disgust, 3:Happiness, 4:Sadness, 5:Anger, 6:Neutral\n",
        "      file_names = train_txtfile.iloc[:, NAME_COLUMN].values\n",
        "    else:\n",
        "      self.label = test_txtfile.iloc[:, LABEL_COLUMN].values - 1 #same as above\n",
        "      file_names = test_txtfile.iloc[:, NAME_COLUMN].values\n",
        "\n",
        "    #Changing file names to match actual image names\n",
        "    self.unmasked_file_paths = []\n",
        "    self.masked_file_paths = []\n",
        "    self.dataset = []\n",
        "    if UU == False and MM == False:\n",
        "      aligned1 = '/aligned'\n",
        "      aligned2 = '/aligned_mask'\n",
        "      file1 = self.root\n",
        "      file2 = self.mask_file\n",
        "    elif UU == True:\n",
        "      aligned1 = '/aligned'\n",
        "      aligned2 = '/aligned'\n",
        "      file1 = self.root\n",
        "      file2 = self.root\n",
        "    elif MM == True:\n",
        "      aligned1 = '/aligned_mask'\n",
        "      aligned2 = '/aligned_mask'\n",
        "      file1 = self.mask_file\n",
        "      file2 = self.mask_file\n",
        "    for f in file_names:\n",
        "        f = f.split(\".\")[0]\n",
        "        f = f +\"_aligned.jpg\"\n",
        "        working_directory = file1 + aligned1\n",
        "        #Putting non-masked image paths into unmasked_file_paths\n",
        "        upath = os.path.join(working_directory, f)\n",
        "        self.unmasked_file_paths.append(upath)\n",
        "        #Putting masked image paths into masked_file_paths\n",
        "        working_directory = file2 + aligned2\n",
        "        mpath = os.path.join(working_directory, f)\n",
        "        self.masked_file_paths.append(mpath)\n",
        "        item = (upath, mpath)\n",
        "        self.dataset.append(item)\n",
        "\n",
        "  def __getitem__(self, index):\n",
        "    \"\"\"\n",
        "    Here we read the actual images\n",
        "    We randomly apply few transforms to the image for image augmentation\n",
        "    Return: image1, image 2 and their label\n",
        "    \"\"\"\n",
        "    \"\"\"\n",
        "    u_imgPath = self.unmasked_file_paths[index]\n",
        "    uimg = cv2.imread(u_imgPath)\n",
        "    uimg = cv2.cvtColor(uimg, cv2.COLOR_BGR2RGB)\n",
        " \n",
        "    m_imgPath = self.masked_file_paths[index]\n",
        "    mimg = cv2.imread(m_imgPath)\n",
        "    mimg = cv2.cvtColor(mimg, cv2.COLOR_BGR2RGB)\n",
        "    \"\"\"\n",
        "    uimg_path, mimg_path = self.dataset[index]\n",
        "    uimg = cv2.imread(uimg_path)\n",
        "    uimg = cv2.cvtColor(uimg, cv2.COLOR_BGR2RGB)\n",
        "    mimg = cv2.imread(mimg_path)\n",
        "    mimg = cv2.cvtColor(mimg, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "\n",
        "    \"\"\"\n",
        "    Note on cv2.imread(): OpenCV uses the BGR format\n",
        "    We need to change to RGB format(ideally before returning the image)\n",
        "    \"\"\"\n",
        "    if self.partition == 'train':\n",
        "      transformed = self.transform1(image = uimg, image1 = mimg)\n",
        "      uimg  = transformed['image']\n",
        "      mimg = transformed['image1']\n",
        "      uimg = self.transform2(uimg)\n",
        "      mimg = self.transform2(mimg)\n",
        "    else:\n",
        "      uimg = self.transform(uimg)\n",
        "      mimg = self.transform(mimg)\n",
        "    label = self.label[index]\n",
        "    return uimg, mimg, label\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.dataset)\n",
        "  \n",
        "  def show_img(self, index):\n",
        "    uimg, mimg, label = self.__getitem__(index)\n",
        "    uimg = uimg.permute(1, 2, 0)\n",
        "    mimg = mimg.permute(1, 2, 0)\n",
        "\n",
        "    f = plt.figure()\n",
        "    f.add_subplot(1,2, 1)\n",
        "    plt.imshow(np.rot90(uimg,0))\n",
        "    f.add_subplot(1,2, 2)\n",
        "    plt.imshow(np.rot90(mimg,0))\n",
        "    plt.show(block=True)\n",
        "\n",
        "    print(\"Label: \", label)\n",
        "\n",
        "#End of Class\n",
        "\n"
      ]
    }
  ]
}