{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch.multiprocessing\n",
    "torch.multiprocessing.set_sharing_strategy('file_system')\n",
    "import torch  \n",
    "import torchvision\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import torchvision.transforms as transforms \n",
    "import numpy as np\n",
    "from sklearn import *\n",
    "import csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pls refer to https://pytorch.org/docs/stable/data.html\n",
    "class EE4146_Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, csv_path, file_path):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            csv_path (string): csv \n",
    "            img_path (string): \n",
    "            transform: transform \n",
    "        \"\"\"\n",
    "        self.file_path = file_path\n",
    "        self.to_tensor = transforms.ToTensor() \n",
    "        self.data_info = pd.read_csv(csv_path, header=0) \n",
    "        self.image_arr = np.asarray(data_info.iloc[:, 0])\n",
    "        self.label_arr = np.asarray(self.data_info.iloc[:, 1])\n",
    "        self.data_len = len(self.data_info.index)\n",
    "        self.encode_labels = dict( zip(np.unique(self.label_arr), range(len(np.unique(self.label_arr)))))\n",
    "        print('encoded labels: ', self.encode_labels)\n",
    "        self.transform = torchvision.transforms.Compose([\n",
    "                    transforms.Resize((150,150)),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize((0.425, 0.415, 0.405), (0.255, 0.245, 0.235))\n",
    "                    ])\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img_id = self.image_arr[index]\n",
    "        img_as_img = Image.open(self.file_path + str(img_id) + \".jpg\")\n",
    "        img_as_img = self.transform(img_as_img)\n",
    "        label = self.encode_labels[self.label_arr[index]]\n",
    "        return (img_as_img, label, img_id)  \n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.data_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define pretrained deep feature extractor\n",
    "model = torchvision.models.wide_resnet50_2(pretrained=True).cuda()\n",
    "for param in model.parameters():\n",
    "    param.required_grad = False\n",
    "\n",
    "# Replacing the final classifier by identity layer (for directly out features)\n",
    "model.fc = torch.nn.Identity()\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoded labels:  {'buildings': 0, 'forest': 1, 'glacier': 2, 'mountain': 3, 'sea': 4, 'street': 5}\n"
     ]
    }
   ],
   "source": [
    "# Define dataloaded\n",
    "train_data = EE4146_Dataset('train_labels.csv','raw_images/train/' )\n",
    "train_loader = torch.utils.data.DataLoader(train_data,batch_size=8,num_workers=0, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference and save features\n",
    "label_save = []\n",
    "feats_save = []\n",
    "feat_save_path = 'train_feat/'\n",
    "\n",
    "if not os.path.exists(feat_save_path):\n",
    "    os.makedirs(feat_save_path)\n",
    "\n",
    "for images, labels, image_ids in train_loader:\n",
    "    B, C, W, H = images.shape\n",
    "    feat = model(images.cuda()).detach().cpu()\n",
    "    feats_save.append(feat) # B 2024\n",
    "    label_save.append(labels)\n",
    "    for i, imgID in enumerate(image_ids):\n",
    "        np.save(feat_save_path +'{}_feat_res50.npy'.format(imgID), feat[i].numpy())\n",
    "\n",
    "feats_save = torch.cat(feats_save,dim=0)\n",
    "label_save = torch.cat(label_save)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "2a1694dcd259e028945710732ec0a65dcf1ed8e1a82793638f5a5001a82d902c"
  },
  "kernelspec": {
   "display_name": "Python 3.6.8 64-bit",
   "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.9.6"
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}