run.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74c60616-08c2-43a4-a548-413deda470b9",
   "metadata": {},
   "source": [
    "# PatchCore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818c359d-add6-47d0-ac70-e724dba21a62",
   "metadata": {},
   "outputs": [],
   "source": [
    "import common\n",
    "import sampler\n",
    "import patchcore\n",
    "import backbones\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a278b2f1-5132-4a4f-9d10-f035e85408f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "# %matplotlib inline\n",
    "\n",
    "from sklearn.model_selection import KFold\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision.transforms import v2\n",
    "\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from argparse import ArgumentParser\n",
    "\n",
    "parser = ArgumentParser(description=\"patchcore\")\n",
    "parser.add_argument('--image_size', default=224, type=int)\n",
    "parser.add_argument('--backbone', default='wideresnet101', type=str)\n",
    "parser.add_argument('--layers_to_extract_from', nargs='+', default=['layer3'], type=str)\n",
    "parser.add_argument('--pretrain_embed_dimension', default=1024, type=int)\n",
    "parser.add_argument('--target_embed_dimension', default=1024, type=int)\n",
    "parser.add_argument('--patchsize', default=3, type=int)\n",
    "parser.add_argument('--anomaly_scorer_num_nn', default=5, type=int)\n",
    "parser.add_argument('--batch_size', default=32, type=int)\n",
    "parser.add_argument('--cv', default=5, type=int)\n",
    "parser.add_argument('--seed', default=826, type=int)\n",
    "parser.add_argument('--device', nargs='+', default=[0], type=int)\n",
    "parser.add_argument('--num_workers', default=0, type=int)\n",
    "args = parser.parse_args('')\n",
    "\n",
    "image_size = args.image_size\n",
    "BATCH_SIZE = args.batch_size\n",
    "CV = args.cv\n",
    "SEED = args.seed\n",
    "\n",
    "def set_seeds(seed=SEED):\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    pl.seed_everything(SEED)\n",
    "\n",
    "set_seeds()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58a005f4-a519-4d1e-9e39-02ff63b4f735",
   "metadata": {},
   "source": [
    "## data_loader.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c637f7-bf1b-4c2e-8fc2-339f50a0210a",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = []\n",
    "y = []\n",
    "\n",
    "for image_path in tqdm(glob(\"data/train/*\")):\n",
    "    X.append(image_path)\n",
    "    y.append(0)\n",
    "    \n",
    "X_test = []\n",
    "for image_path in tqdm(glob(\"data/test/*\")):\n",
    "    X_test.append(image_path)\n",
    "\n",
    "X = np.array(X)\n",
    "y = np.array(y)\n",
    "X_test = np.array(X_test)\n",
    "\n",
    "len(X), len(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f194bc-f437-4bc2-b33f-8c1d5a9bcd5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransistorDataset(Dataset):\n",
    "    def __init__(self, X, y=None, transform=None):\n",
    "        super().__init__()        \n",
    "        self.X = X\n",
    "        self.y = y\n",
    "        self.transform = transform\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        X = self.X[idx]\n",
    "        X = Image.open(X)\n",
    "        \n",
    "        if self.transform is not None:\n",
    "            X = self.transform(X)\n",
    "        \n",
    "        if self.y is None:\n",
    "            return {\"image\": X}\n",
    "        \n",
    "        y = self.y[idx]        \n",
    "        return {\"image\": X, \"label\": y}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60c4597-b6e5-48b6-9bb7-a20ede34c811",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transform = v2.Compose([\n",
    "    v2.ToImage(),\n",
    "    v2.Resize(size=(round(image_size*1.143), round(image_size*1.143))),\n",
    "    v2.CenterCrop(size=(image_size, image_size)),\n",
    "    v2.ToDtype(torch.float32, scale=True),\n",
    "    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "\n",
    "val_transform = v2.Compose([\n",
    "    v2.ToImage(),\n",
    "    v2.Resize(size=(round(image_size*1.143), round(image_size*1.143))),\n",
    "    v2.CenterCrop(size=(image_size, image_size)),\n",
    "    v2.ToDtype(torch.float32, scale=True),\n",
    "    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "\n",
    "test_transform = v2.Compose([\n",
    "    v2.ToImage(),\n",
    "    v2.Resize(size=(round(image_size*1.143), round(image_size*1.143))),\n",
    "    v2.CenterCrop(size=(image_size, image_size)),\n",
    "    v2.ToDtype(torch.float32, scale=True),\n",
    "    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2523acc-371a-467b-86ba-b93381084afb",
   "metadata": {},
   "source": [
    "## train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff663d48-fcb7-4577-9f39-110a565f6c03",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "threshold_list = [] \n",
    "scores_list = [] \n",
    "\n",
    "kf = KFold(n_splits=CV, shuffle=True, random_state=SEED)\n",
    "\n",
    "for i, (train_index, val_index) in enumerate(kf.split(X)):\n",
    "\n",
    "    X_train = X[train_index]\n",
    "    y_train = y[train_index]\n",
    "    X_val = X[val_index]\n",
    "    y_val = y[val_index]\n",
    "\n",
    "    train_dataset = TransistorDataset(X_train, y_train, train_transform)\n",
    "    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=args.num_workers)\n",
    "    \n",
    "    val_dataset = TransistorDataset(X_val, y_val, val_transform)\n",
    "    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=args.num_workers)\n",
    "    \n",
    "    test_dataset = TransistorDataset(X_test, None, test_transform)\n",
    "    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=args.num_workers)\n",
    "\n",
    "    device = utils.set_torch_device(gpu_ids=args.device)\n",
    "    \n",
    "    patch_core = patchcore.PatchCore(device)\n",
    "    \n",
    "    patch_core.load(\n",
    "        backbone                 = backbones.load(args.backbone),\n",
    "        layers_to_extract_from   = args.layers_to_extract_from,\n",
    "        device                   = device,\n",
    "        input_shape              = (3, image_size, image_size),\n",
    "        pretrain_embed_dimension = args.pretrain_embed_dimension,\n",
    "        target_embed_dimension   = args.target_embed_dimension,\n",
    "        patchsize                = args.patchsize,\n",
    "        anomaly_scorer_num_nn    = args.anomaly_scorer_num_nn,\n",
    "        featuresampler           = sampler.GreedyCoresetSampler(percentage=0.1, device=device),\n",
    "        nn_method                = common.FaissNN(on_gpu=False, num_workers=args.num_workers)\n",
    "    )\n",
    "\n",
    "    patch_core.fit(train_dataloader)\n",
    "\n",
    "    scores, _ = patch_core.predict(\n",
    "        val_dataloader\n",
    "    )\n",
    "\n",
    "    threshold = np.max(scores)\n",
    "    print(f\"threshold: {threshold}\")\n",
    "\n",
    "    scores, _ = patch_core.predict(\n",
    "        test_dataloader\n",
    "    )\n",
    "\n",
    "    threshold_list.append(threshold)\n",
    "    scores_list.append(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4796ce0b-353b-484a-ba4b-034e40ae6e23",
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = np.mean(threshold_list)\n",
    "print(f\"threshold: {threshold}\")\n",
    "\n",
    "scores = np.max(scores_list, axis=0)\n",
    "prediction = np.where(scores<threshold, 0, 1)\n",
    "\n",
    "print(f\"n_anomaly: {np.sum(prediction)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3671dc74-3013-4827-ab98-a446e2c013e6",
   "metadata": {},
   "source": [
    "## submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106ad99b-800e-4c33-a066-95bf9d93f7d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_submission = pd.read_csv('data/sample_submission.csv')\n",
    "sample_submission['label'] = prediction\n",
    "sample_submission.to_csv('patchcore.csv', index=False)\n",
    "\n",
    "sample_submission.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}