create_finetune_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader, Subset\n",
    "from torch.optim import AdamW\n",
    "from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, fasterrcnn_resnet50_fpn\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.ops import nms\n",
    "from torchvision.io import read_image\n",
    "import torchvision.transforms.v2 as T\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import utils1\n",
    "from utils1 import * \n",
    "from evaluation import evaluate_with_results\n",
    "from datasets import *\n",
    "from torchvision.ops import nms, box_iou\n",
    "import torchvision.transforms as transforms\n",
    "from visualisation import *\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "#Settings\n",
    "images_path_GTSDB = \"GTSDBDataset\"\n",
    "annotaions_path_GTSDB = \"GTSDBDataset/gt.txt\"\n",
    "localization_model_path = \"Models\\modelGTSDB_WithoutClasses19.pth\"\n",
    "new_dataset_path = \"FineTuneDataset\"\n",
    "batch_size = 4\n",
    "nms_threshold = 0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "localization_model = torch.load(localization_model_path)\n",
    "localization_model.to(device)\n",
    "localization_model.eval()\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_GTSDB = GTSDBDataset(True, images_path_GTSDB, annotaions_path_GTSDB)\n",
    "dataloader_GTSDB = DataLoader(dataset_GTSDB, batch_size=batch_size, pin_memory=True, shuffle=False, collate_fn=collate_fn)\n",
    "\n",
    "train_indices = list(range(600))\n",
    "train_dataset = Subset(dataset_GTSDB, train_indices)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, pin_memory=True, shuffle=False, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_matching_box(boxes, target_box):\n",
    "    target_box = target_box.unsqueeze(0)\n",
    "    ious = box_iou(boxes, target_box)\n",
    "    if ious.size(0) == 0:\n",
    "        return 0, 0\n",
    "    iou, idx = ious.max(0)\n",
    "    return iou, idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes_images = {}\n",
    "for i in range(44):\n",
    "    classes_images[i] = []\n",
    "\n",
    "for data in tqdm(train_dataloader):\n",
    "    for i in range(batch_size):\n",
    "        image = data[0][i].to(device)\n",
    "        image_to_show = cv2.imread(data[2][i])\n",
    "\n",
    "        outputs = localization_model([image])\n",
    "        boxes = outputs[0][\"boxes\"].detach().cpu()\n",
    "        scores = outputs[0][\"scores\"].detach().cpu()\n",
    "        labels = outputs[0][\"labels\"].detach().cpu()\n",
    "\n",
    "        to_keep = nms(boxes, scores, nms_threshold)\n",
    "        boxes = boxes[to_keep]\n",
    "        scores = scores[to_keep]\n",
    "\n",
    "        images = extract_signs(image_to_show, boxes)\n",
    "        actual = {}\n",
    "        for i2, box in enumerate(data[1][i][\"boxes\"]):\n",
    "            actual[box] = data[1][i][\"labels\"][i2].item()\n",
    "\n",
    "        for i2 in range(len(images)):\n",
    "            box = boxes[i2]\n",
    "            image = images[i2]\n",
    "            iou, idx = find_matching_box(data[1][i][\"boxes\"], box)\n",
    "            if iou < 0.7:\n",
    "                #wrong box detected\n",
    "                label = 0 #0 = background / none\n",
    "            else:\n",
    "                label = data[1][i][\"labels\"][idx].item()\n",
    "            classes_images[label].append(image)\n",
    "       \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(44):\n",
    "    print(\"Class \" + class_map[i] + \" has \" + str(len(classes_images[i])) + \" images\")\n",
    "    for i2 in range(len(classes_images[i])):\n",
    "        if not os.path.exists(new_dataset_path + \"/\" + str(i)):\n",
    "            os.makedirs(new_dataset_path +\"/\" + str(i))\n",
    "        cv2.imwrite(new_dataset_path + \"/\" + str(i) + \"/\" + str(i2) + \".jpg\", classes_images[i][i2])\n",
    "        # cv2.imshow(\"Class \" + str(i), classes_images[i][i2])\n",
    "        # cv2.waitKey(0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fm",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}