training_vehicle_reid.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vehicle Re-Identification Training Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mdxuser/anaconda3/envs/carla_reid/lib/python3.7/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torchreid\n",
    "import sys\n",
    "import os.path as osp\n",
    "import pandas as pd\n",
    "from torchreid.data import ImageDataset\n",
    "import torch\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torchreid versions is: 1.4.0\n"
     ]
    }
   ],
   "source": [
    "print(f\"torchreid versions is: {torchreid.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n",
    "os.makedirs(\"./log\", exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some of the available models, which you can directly use with torchreid module\n",
    "models_list = ['osnet_x1_0' ,'osnet_ibn_x1_0', 'osnet_ain_x1_0', 'resnet50','resnet152']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VeRiDataset(ImageDataset):\n",
    "\n",
    "    # All you need to do here is to generate three lists,\n",
    "    # which are train, query and gallery.\n",
    "    # Each list contains tuples of (img_path, pid, camid),\n",
    "    # where\n",
    "    # - img_path (str): absolute path to an image.\n",
    "    # - pid (int): person ID, e.g. 0, 1.\n",
    "    # - camid (int): camera ID, e.g. 0, 1.\n",
    "    # Note that\n",
    "    # - pid and camid should be 0-based.\n",
    "    # - query and gallery should share the same pid scope (e.g.\n",
    "    #   pid=0 in query refers to the same person as pid=0 in gallery).  --> Very important!\n",
    "    # - train, query and gallery share the same camid scope (e.g.\n",
    "    #   camid=0 in train refers to the same camera as camid=0\n",
    "    #   in query/gallery).\n",
    "\n",
    "    def __init__(self, root='', **kwargs):\n",
    "        PATH_TO_DATASET = \"/home/mdxuser/Desktop/Ph.D/GitHub/VehicleReIdentificationDataset/dataset/\" # This is the folder path where you have extracted the zip file and it contains the gallery, query and train images\n",
    "        FOLDER_PATH, DATASET_TYPES, IMAGES = next(os.walk(f\"{PATH_TO_DATASET}\")) \n",
    "        training_array = []\n",
    "        testing_array = []\n",
    "        query_array = []\n",
    "\n",
    "        for dataset_type in DATASET_TYPES:\n",
    "            _, _, IMAGES = next(os.walk(f\"{FOLDER_PATH}/{dataset_type}/\"))\n",
    "\n",
    "            if dataset_type == \"image_train\" or dataset_type == \"image_gallery\":\n",
    "\n",
    "                labelSet = []\n",
    "                for idx, image in enumerate(IMAGES):\n",
    "                    pid = int(image.split('.jpg')[0].split('_')[-1])\n",
    "                    labelSet.append(pid)\n",
    "                label_dict = {label: index for index, label in enumerate(set(labelSet))}\n",
    "                \n",
    "                for image in IMAGES:\n",
    "                    cam_id = int(image.split('.jpg')[0].split('_')[-2])\n",
    "                    car_id = label_dict[int(image.split('.jpg')[0].split('_')[-1])]\n",
    "\n",
    "                    if dataset_type == \"image_gallery\":\n",
    "                        testing_array.append((FOLDER_PATH +  dataset_type + \"/\" + image, car_id, cam_id))\n",
    "                    else:\n",
    "                        training_array.append((FOLDER_PATH +  dataset_type + \"/\" + image, car_id, cam_id))\n",
    "                        \n",
    "            elif dataset_type == \"image_query\":\n",
    "                for image in IMAGES:\n",
    "                    cam_id = int(image.split('.jpg')[0].split('_')[-2])\n",
    "                    car_id = label_dict[int(image.split('.jpg')[0].split('_')[-1])]\n",
    "                    query_array.append((FOLDER_PATH +  dataset_type + \"/\" + image, car_id, cam_id))\n",
    "\n",
    "            \n",
    "        train = training_array\n",
    "        query = query_array\n",
    "        gallery = testing_array\n",
    "\n",
    "        super(VeRiDataset, self).__init__(train, query, gallery, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    torchreid.data.register_image_dataset('synthtetic_veri_reid', VeRiDataset)\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building train transforms ...\n",
      "+ resize to 256x256\n",
      "+ random flip\n",
      "+ random crop (enlarge to 288x288 and crop 256x256)\n",
      "+ to torch tensor of range [0, 1]\n",
      "+ normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "Building test transforms ...\n",
      "+ resize to 256x256\n",
      "+ to torch tensor of range [0, 1]\n",
      "+ normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "=> Loading train (source) dataset\n",
      "=> Loaded VeRiDataset\n",
      "  ----------------------------------------\n",
      "  subset   | # ids | # images | # cameras\n",
      "  ----------------------------------------\n",
      "  train    |   600 |    50949 |        85\n",
      "  query    |    50 |      424 |        85\n",
      "  gallery  |    50 |     3823 |        85\n",
      "  ----------------------------------------\n",
      "=> Loading test (target) dataset\n",
      "=> Loaded VeRiDataset\n",
      "  ----------------------------------------\n",
      "  subset   | # ids | # images | # cameras\n",
      "  ----------------------------------------\n",
      "  train    |   600 |    50949 |        85\n",
      "  query    |    50 |      424 |        85\n",
      "  gallery  |    50 |     3823 |        85\n",
      "  ----------------------------------------\n",
      "\n",
      "\n",
      "  **************** Summary ****************\n",
      "  source            : ['synthtetic_veri_reid']\n",
      "  # source datasets : 1\n",
      "  # source ids      : 600\n",
      "  # source images   : 50949\n",
      "  # source cameras  : 85\n",
      "  target            : ['synthtetic_veri_reid']\n",
      "  *****************************************\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# You can change these hyperparameters as you wish\n",
    "datamanager = torchreid.data.ImageDataManager(\n",
    "    root='',\n",
    "    sources='synthtetic_veri_reid',\n",
    "    height = 256,\n",
    "    width = 256,\n",
    "    transforms =['random_flip', 'random_crop'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOSS = \"SOFTMAX\" # Define softmax loss, else triple loss will be usd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "osnet_x1_0\n",
      "Successfully loaded imagenet pretrained weights from \"/home/mdxuser/.cache/torch/checkpoints/osnet_x1_0_imagenet.pth\"\n",
      "** The following layers are discarded due to unmatched keys or layer size: ['classifier.weight', 'classifier.bias']\n"
     ]
    }
   ],
   "source": [
    "for current_model in models_list:\n",
    "    print(current_model)\n",
    "    if LOSS == \"SOFTMAX\":\n",
    "        model = torchreid.models.build_model(\n",
    "            name=f'{current_model}',\n",
    "            num_classes=datamanager.num_train_pids,\n",
    "            loss='softmax'\n",
    "        )\n",
    "    else:\n",
    "        model = torchreid.models.build_model(\n",
    "            name=f'{current_model}',\n",
    "            num_classes=datamanager.num_train_pids,\n",
    "            loss='triplet'\n",
    "        )\n",
    "\n",
    "    model = model.cuda()\n",
    "    optimizer = torchreid.optim.build_optimizer(\n",
    "        model,\n",
    "        optim='adam',\n",
    "        lr=0.0003\n",
    "    )\n",
    "\n",
    "    scheduler = torchreid.optim.build_lr_scheduler(\n",
    "        optimizer,\n",
    "        lr_scheduler='single_step',\n",
    "        stepsize=20\n",
    "    )\n",
    "\n",
    "    if LOSS == \"SOFTMAX\":\n",
    "        engine = torchreid.engine.ImageSoftmaxEngine(\n",
    "            datamanager,\n",
    "            model,\n",
    "            optimizer=optimizer,\n",
    "            scheduler=scheduler,\n",
    "            \n",
    "        )\n",
    "    else:\n",
    "        print(\"Using ImageTripletEngine\")\n",
    "        engine = torchreid.engine.ImageTripletEngine(\n",
    "            datamanager,\n",
    "            model,\n",
    "            optimizer=optimizer,\n",
    "            scheduler=scheduler,\n",
    "        )\n",
    "\n",
    "    original_stdout = sys.stdout\n",
    "    with open(f\"log/log_{current_model}.txt\", 'a') as f:\n",
    "        sys.stdout = f\n",
    "        engine.run(\n",
    "        save_dir=f\"log/{current_model}\",\n",
    "        max_epoch=60,\n",
    "        eval_freq=5,\n",
    "        print_freq=10,\n",
    "        test_only=False\n",
    "        )\n",
    "        sys.stdout = original_stdout\n",
    "\n",
    "    del engine\n",
    "    del model\n",
    "    del optimizer\n",
    "    del scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "carla_reid",
   "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.7.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8d00adedad053f31af55c36cf95aece13fe582b912fc1bbfa923bebe8f098642"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}