viz_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from data import StratifiedGroupKFoldDataModule\n",
    "import cv2 as cv\n",
    "from pathlib import Path\n",
    "from models.base import BaseModel\n",
    "from models.resnet import ResNetModel\n",
    "from torch.nn import BCEWithLogitsLoss\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "group = parser.add_argument_group(\"qt.coyote\")\n",
    "group.add_argument(\"--batch_size\", help=\"Batch size\", type=int, default=32)\n",
    "group.add_argument(\n",
    "    \"--learning_rate\", help=\"Learning rate\", type=float, default=1e-3\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--k\",\n",
    "    help=\"Number of folds in k-fold cross validation\",\n",
    "    type=int,\n",
    "    default=5,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--internal_k\",\n",
    "    help=\"Number of folds for train/test split\",\n",
    "    type=int,\n",
    "    default=5,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--no_external_group\",\n",
    "    help=\"Use grouped k-fold cross validation in external k-fold cross validation\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--internal_group\",\n",
    "    help=\"Use grouped k-fold cross validation in internal k-fold cross validation\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--data_path\",\n",
    "    help=\"Path to images\",\n",
    "    type=str,\n",
    "    default=\"data\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--metadata_path\",\n",
    "    help=\"Path to COCO metadata file\",\n",
    "    type=str,\n",
    "    default=\"data/qt-coyotes-merged.json\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--num_workers\",\n",
    "    help=\"Number of workers for dataloader\",\n",
    "    type=int,\n",
    "    default=os.cpu_count() - 2,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--persistent_workers\",\n",
    "    help=\"If True, the data loader will not shutdown the worker processes \"\n",
    "    \"after a dataset has been consumed once. This allows to maintain the \"\n",
    "    \"workers Dataset instances alive.\",\n",
    "    type=bool,\n",
    "    default=True,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--shuffle\",\n",
    "    help=\"Whether to shuffle each class's samples before splitting into \"\n",
    "    \"batches. Note that the samples within each split will not be \"\n",
    "    \"shuffled. This implementation can only shuffle groups that have \"\n",
    "    \"approximately the same y distribution, no global shuffle will be \"\n",
    "    \"performed.\",\n",
    "    type=bool,\n",
    "    default=True,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--random_state\",\n",
    "    help=\"When shuffle is True, random_state affects the ordering of the \"\n",
    "    \"indices, which controls the randomness of each fold for each class. \"\n",
    "    \"Otherwise, leave random_state as None. Pass an int for reproducible \"\n",
    "    \"output across multiple function calls.\",\n",
    "    type=int,\n",
    "    default=42,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--nondeterministic\",\n",
    "    help=\"This flag sets the torch.backends.cudnn.deterministic flag to false\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--nonpretrained\",\n",
    "    help=\"Do not use pretrained weights, train from scratch\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--compile\",\n",
    "    help=\"Compile the model\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--no_early_stopping\",\n",
    "    help=\"Disable early stopping\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--no_crop\",\n",
    "    help=\"Disable cropping\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--no_data_augmentation\",\n",
    "    help=\"Disable data augmentation\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--patience\",\n",
    "    help=\"Number of checks with no improvement after which training will be stopped.\",\n",
    "    type=int,\n",
    "    default=5,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--scheduler_factor\",\n",
    "    help=\"Factor by which the lr will be decreased\",\n",
    "    type=float,\n",
    "    default=0.5,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--scheduler_patience\",\n",
    "    help=\"Number of checks with no improvement after which lr will decrease\",\n",
    "    type=int,\n",
    "    default=4,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--yolo_model\",\n",
    "    help=\"Yolo pretrained model\",\n",
    "    type=str,\n",
    "    default=\"yolov8n-cls.pt\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--resnet_model\",\n",
    "    help=\"ResNet model\",\n",
    "    type=str,\n",
    "    default=\"ResNet18\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--densenet_model\",\n",
    "    help=\"DenseNet model\",\n",
    "    type=str,\n",
    "    default=\"DenseNet121\",\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--crop_size\",\n",
    "    help=\"Crop size\",\n",
    "    type=int,\n",
    "    default=224,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--criterion_pos_weight\",\n",
    "    help=\"Weight for positive class for BCEWithLogitsLoss\",\n",
    "    type=float,\n",
    "    default=10.0\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--criterion_beta\",\n",
    "    help=\"Beta for F-beta loss\",\n",
    "    type=float,\n",
    "    default=5.0\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--criterion_cfn\",\n",
    "    help=\"Cost false negative for ExpectedCostLoss\",\n",
    "    type=float,\n",
    "    default=5.0\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--dropout_p\",\n",
    "    help=\"Dropout probability\",\n",
    "    type=float,\n",
    "    default=0.2,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--tabular_hidden_size\",\n",
    "    help=\"Size of the tabular hidden layers\",\n",
    "    type=int,\n",
    "    default=32,\n",
    ")\n",
    "group.add_argument(\n",
    "    \"--no_tabular_features\",\n",
    "    help=\"Do not use tabular features\",\n",
    "    action=\"store_true\",\n",
    ")\n",
    "args = parser.parse_args(\"--batch_size 16 --learning_rate 0.0001 --patience 5\".split())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as transforms\n",
    "augmentations = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomRotation(10),\n",
    "    transforms.GaussianBlur(3, sigma=(0.1, 2)),\n",
    "    transforms.ColorJitter(\n",
    "        brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1\n",
    "    ),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = BCEWithLogitsLoss()\n",
    "model = ResNetModel(loss, args)\n",
    "print(model.feature_extractor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test for external k-fold cross validation leakage of groups\n",
    "datamodule = StratifiedGroupKFoldDataModule(args)\n",
    "datamodule.prepare_data()\n",
    "datamodule.setup(None)\n",
    "\n",
    "locations = {}\n",
    "\n",
    "for i, datamodule_i in tqdm(enumerate(datamodule)):\n",
    "    train = datamodule_i.train_dataloader()\n",
    "    val = datamodule_i.val_dataloader()\n",
    "    test = datamodule_i.test_dataloader()\n",
    "    train_locations = train.dataset.locations\n",
    "    val_locations = val.dataset.locations\n",
    "    test_locations = test.dataset.locations\n",
    "    locations[i] = {\n",
    "        \"train\": list(train_locations),\n",
    "        \"val\": list(val_locations),\n",
    "        \"test\": list(test_locations),\n",
    "    }\n",
    "    assert(test_locations.intersection(train_locations) == set())\n",
    "    assert(test_locations.intersection(val_locations) == set())\n",
    "\n",
    "\n",
    "import json\n",
    "with open(\"locations.json\", \"w\") as f:\n",
    "    json.dump(locations, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.dataset.locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "datamodule = StratifiedGroupKFoldDataModule(args)\n",
    "datamodule.prepare_data()\n",
    "datamodule.setup(None)\n",
    "\n",
    "DEBUG_PATH = Path(\"data/debug\")\n",
    "\n",
    "criterion = BCEWithLogitsLoss()\n",
    "\n",
    "import math\n",
    "for i, datamodule_i in tqdm(enumerate(datamodule)):\n",
    "    train = datamodule_i.train_dataloader()\n",
    "    val = datamodule_i.val_dataloader()\n",
    "    test = datamodule_i.test_dataloader()\n",
    "    ltrain = len(train.dataset)\n",
    "    lval = len(val.dataset)\n",
    "    ltest = len(test.dataset)\n",
    "    print(f\"Fold {i}: train={ltrain}, val={lval}, test={ltest}, total={ltrain+lval+ltest}\")\n",
    "    folder_path = DEBUG_PATH / f\"fold_{i}\"\n",
    "    model_path = glob.glob(f\"lightning_logs/version_{i}/checkpoints/*.ckpt\")[0]\n",
    "    model = ResNetModel(criterion, args)\n",
    "    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))[\"state_dict\"], strict=False)\n",
    "    model.eval()\n",
    "    # for dataloader, stage in zip([train, val, test], [\"train\", \"val\", \"test\"]):\n",
    "    for dataloader, stage in zip([test], [\"test\"]):\n",
    "        stage_path = folder_path / stage\n",
    "\n",
    "        for j, (X, Y) in tqdm(enumerate(dataloader)):\n",
    "            fig, axs = plt.subplots(4, 4, figsize=(16, 16))\n",
    "            if stage == \"test\":\n",
    "                with torch.no_grad():\n",
    "                    LOGITS = model.y(X)\n",
    "                    YHAT = torch.sigmoid(LOGITS)\n",
    "            else:\n",
    "                YHAT = [None] * len(Y)\n",
    "            for k, (i, t, y, yhat) in enumerate(zip(*X, Y, YHAT)):\n",
    "                image_path = stage_path / f\"{y}/batch_{j}_{k}.png\"\n",
    "                image_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "                ax = axs[k // 4, k % 4]\n",
    "                i = i.permute(1, 2, 0)\n",
    "                cv.imwrite(str(image_path), i.numpy())\n",
    "                ax.imshow(i)\n",
    "                    \n",
    "                if stage == \"test\":\n",
    "                    ax.set_title(f\"{stage} {yhat:.2f}: {y}\")\n",
    "                else:\n",
    "                    ax.set_title(f\"{stage} {y}\")\n",
    "            batch_path = stage_path / f\"batch_{j}.png\"\n",
    "            batch_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "            plt.savefig(batch_path)\n",
    "            plt.close()\n",
    "\n",
    "            if stage == \"test\":\n",
    "                continue\n",
    "            Xa = augmentations(X[0])\n",
    "            fig, axs = plt.subplots(4, 4, figsize=(16, 16))\n",
    "            for k, (x, y) in enumerate(zip(Xa, Y)):\n",
    "                ax = axs[k // 4, k % 4]\n",
    "                ax.imshow(i.permute(1, 2, 0))\n",
    "                ax.set_title(f\"{stage} {y}\")\n",
    "            batch_path = stage_path / f\"batch_{j}_augmented.png\"\n",
    "            batch_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "            plt.savefig(batch_path)\n",
    "            plt.close()\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "\n",
    "with open(\"data/qt-coyotes-merged.json\") as f:\n",
    "    coco = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = coco[\"images\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "year = np.array([image[\"year\"] for image in images])\n",
    "month = np.array([image[\"month\"] for image in images])\n",
    "day = np.array([image[\"day\"] for image in images])\n",
    "hour = np.array([image[\"hour\"] for image in images])\n",
    "is_color = np.array([int(1 if image[\"is_color\"] else 0) for image in images])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "for a in [year, month, day, hour, is_color]:\n",
    "    print(a.mean(), a.std())\n",
    "    plt.hist(a)\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f80fc6aa78b7af1cf9644414b936d87a8b782647fff8191a8cc6b562d64bd9bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}