notebooks/200_models/201_fastflow.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Working Directory\n",
    "This cell is to ensure we change the directory to anomalib source code to have access to the datasets and config files. We assume that you already went through `001_getting_started.ipynb` and install the required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from functools import partial, update_wrapper\n",
    "from pathlib import Path\n",
    "from types import MethodType\n",
    "from typing import Any\n",
    "\n",
    "from git.repo import Repo\n",
    "\n",
    "current_directory = Path.cwd()\n",
    "if current_directory.name == \"200_models\":\n",
    "    # On the assumption that, the notebook is located in\n",
    "    #   ~/anomalib/notebooks/100_datamodules/\n",
    "    root_directory = current_directory.parent.parent\n",
    "elif current_directory.name == \"anomalib\":\n",
    "    # This means that the notebook is run from the main anomalib directory.\n",
    "    root_directory = current_directory\n",
    "else:\n",
    "    # Otherwise, we'll need to clone the anomalib repo to the `current_directory`\n",
    "    repo = Repo.clone_from(url=\"https://github.com/openvinotoolkit/anomalib.git\", to_path=current_directory)\n",
    "    root_directory = current_directory / \"anomalib\"\n",
    "\n",
    "os.chdir(root_directory)\n",
    "dataset_root = root_directory / \"datasets\" / \"MVTec\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Train a Model via API\n",
    "\n",
    "This notebook demonstrates how to train, test and infer the FastFlow model via Anomalib API. Compared to the CLI entrypoints such as \\`tools/\\<train, test, inference>.py, the API offers more flexibility such as modifying the existing model or designing custom approaches.\n",
    "\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from pytorch_lightning import LightningModule, Trainer\n",
    "from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n",
    "from torch.optim import Optimizer\n",
    "from torch.optim.adam import Adam\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from anomalib.data import InferenceDataset, TaskType\n",
    "from anomalib.data.mvtec import MVTec\n",
    "from anomalib.models.fastflow.lightning_model import Fastflow\n",
    "from anomalib.post_processing import (\n",
    "    NormalizationMethod,\n",
    "    ThresholdMethod,\n",
    "    superimpose_anomaly_map,\n",
    ")\n",
    "from anomalib.pre_processing.transforms import Denormalize\n",
    "from anomalib.utils.callbacks import (\n",
    "    ImageVisualizerCallback,\n",
    "    MetricsConfigurationCallback,\n",
    "    MetricVisualizerCallback,\n",
    "    PostProcessingConfigurationCallback,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Data Module\n",
    "\n",
    "To train the model end-to-end, we do need to have a dataset. In our [previous notebooks](https://github.com/openvinotoolkit/anomalib/tree/main/notebooks/100_datamodules), we demonstrate how to initialize benchmark- and custom datasets. In this tutorial, we will use MVTec AD DataModule. We assume that `datasets` directory is created in the `anomalib` root directory and `MVTec` dataset is located in `datasets` directory.\n",
    "\n",
    "Before creating the dataset, let's define the task type that we will be working on. In this notebook, we will be working on a segmentation task. Therefore the `task` variable would be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = TaskType.SEGMENTATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "datamodule = MVTec(\n",
    "    root=dataset_root,\n",
    "    category=\"bottle\",\n",
    "    image_size=256,\n",
    "    train_batch_size=32,\n",
    "    eval_batch_size=32,\n",
    "    num_workers=8,\n",
    "    task=task,\n",
    ")\n",
    "datamodule.setup()\n",
    "i, data = next(enumerate(datamodule.test_dataloader()))\n",
    "print(f'Image Shape: {data[\"image\"].shape} Mask Shape: {data[\"mask\"].shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## FastFlow Model\n",
    "\n",
    "Now that we have created the MVTec datamodule, we could create the FastFlow model. We could start with printing its docstring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "Fastflow??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model = Fastflow(input_size=(256, 256), backbone=\"resnet18\", flow_steps=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Depending on the `training` mode, `model` returns two different outputs. If the model is in `training` mode, it returns the hidden variable and the log of the jacobian, based on the feature maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model.training = True\n",
    "train_output = model(data[\"image\"])\n",
    "hidden_variables, log_jacobian = train_output\n",
    "print(f\"Hidden Variable Shape: {hidden_variables[0].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "During the test/inference mode, the model returns an anomaly heatmap localizing the anomalous regions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model.model.training = False\n",
    "anomaly_map = model(data[\"image\"])\n",
    "print(f\"Anomaly Map Shape: {anomaly_map.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Optimizer\n",
    "\n",
    "`LightningModule` has `configure_optimizer` method that returns the optimizer object. This is not implemented in the FastFlow model. Instead, it is left to the user to make it configurable by either CLI or API. Here we will be configuring it via the API. To do so, we first need to create the optimizer object, wrap it around a function and finally inject this function into the model class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def configure_optimizers(lightning_module: LightningModule, optimizer: Optimizer) -> Any:  # pylint: disable=W0613,W0621\n",
    "    \"\"\"Override to customize the LightningModule.configure_optimizers` method.\"\"\"\n",
    "    return optimizer\n",
    "\n",
    "\n",
    "optimizer = Adam(params=model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=1e-5)\n",
    "fn = partial(configure_optimizers, optimizer=optimizer)\n",
    "update_wrapper(fn, configure_optimizers)  # necessary for `is_overridden`\n",
    "model.configure_optimizers = MethodType(fn, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Callbacks\n",
    "\n",
    "To train the model properly, we will to add some other \"non-essential\" logic such as saving the weights, early-stopping, normalizing the anomaly scores and visualizing the input/output images. To achieve these we use `Callbacks`. Anomalib has its own callbacks and also supports PyTorch Lightning's native callbacks. So, let's create the list of callbacks we want to execute during the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "callbacks = [\n",
    "    MetricsConfigurationCallback(\n",
    "        task=task,\n",
    "        image_metrics=[\"AUROC\"],\n",
    "        pixel_metrics=[\"AUROC\"],\n",
    "    ),\n",
    "    ModelCheckpoint(\n",
    "        mode=\"max\",\n",
    "        monitor=\"pixel_AUROC\",\n",
    "    ),\n",
    "    EarlyStopping(\n",
    "        monitor=\"pixel_AUROC\",\n",
    "        mode=\"max\",\n",
    "        patience=3,\n",
    "    ),\n",
    "    PostProcessingConfigurationCallback(\n",
    "        normalization_method=NormalizationMethod.MIN_MAX,\n",
    "        threshold_method=ThresholdMethod.ADAPTIVE,\n",
    "    ),\n",
    "    ImageVisualizerCallback(mode=\"full\", task=task, image_save_path=\"./results/images\"),\n",
    "    MetricVisualizerCallback(mode=\"full\", task=task, image_save_path=\"./results/images\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training\n",
    "\n",
    "Now that we set up the datamodule, model, optimizer and the callbacks, we could now train the model.\n",
    "\n",
    "The final component to train the model is `pytorch_lightning` `Trainer` object, which handles train/test/predict pipeline. Let's create the trainer object to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    callbacks=callbacks,\n",
    "    accelerator=\"auto\",  # \\<\"cpu\", \"gpu\", \"tpu\", \"ipu\", \"hpu\", \"auto\">,\n",
    "    devices=1,\n",
    "    max_epochs=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "`Trainer` object has number of options that suit all specific needs. For more details, refer to [Lightning Documentation](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html) to see how it could be tweaked to your needs.\n",
    "\n",
    "Let's train the model now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.fit(datamodule=datamodule, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The training has finished after 12 epochs. This is because, we set the `EarlyStopping` criteria with a patience of 3, which terminated the training after `pixel_AUROC` stopped improving. If we increased the `patience`, the training would continue further.\n",
    "\n",
    "## Testing\n",
    "\n",
    "Now that we trained the model, we could test the model to check the overall performance on the test set. We will also be writing the output of the test images to a file since we set `VisualizerCallback` in `callbacks`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.test(datamodule=datamodule, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "`trainer.test` returns the `pixel_AUROC` and `image_AUROC` results. We could also find the saved output in `images` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ls images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Inference\n",
    "\n",
    "Since we have a trained model, we could infer the model on an individual image or folder of images. Anomalib has an `InferenceDataset` to let you create an inference dataset. So let's try it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "inference_dataset = InferenceDataset(path=dataset_root / \"bottle/test/broken_large/000.png\", image_size=(256, 256))\n",
    "inference_dataloader = DataLoader(dataset=inference_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We could utilize `Trainer`'s `predict` method to infer, and get the outputs to visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "predictions = trainer.predict(model=model, dataloaders=inference_dataloader)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "`predictions` contain image, anomaly maps, predicted scores, labels and masks. These are all stored in a dictionary. We could check this by printing the `prediction` keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(predictions.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(\n",
    "    f'Image Shape: {predictions[\"image\"].shape},\\n'\n",
    "    'Anomaly Map Shape: {predictions[\"anomaly_maps\"].shape}, \\n'\n",
    "    'Predicted Mask Shape: {predictions[\"pred_masks\"].shape}'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To properly visualize the predictions, we will need to perform some post-processing operations.\n",
    "\n",
    "Let's post-process each output one by one. We could start with the image. Each `image` is a tensor and within (0, 1) range. To visualize it, we need to denormalize it to (0, 255) scale. Anomalib already has a class for this. Let's use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "image = predictions[\"image\"][0]\n",
    "image = Denormalize()(image)\n",
    "print(f\"Image Shape: {image.shape}\\n Min Pixel: {image.min()} \\n Max Pixel: {image.max()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We could now see that `image` is of a numpy array and within the range of \\[0, 255\\]. It's ready to be visualized now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "plt.imshow(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The second output of the predictions is the anomaly map. As can be seen above, it's also a torch tensor and of size `torch.Size([1, 1, 256, 256])`. We therefore need to convert it to numpy and squeeze the dimensions to make it `256x256` output to visualize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "anomaly_map = predictions[\"anomaly_maps\"][0]\n",
    "anomaly_map = anomaly_map.cpu().numpy().squeeze()\n",
    "plt.imshow(anomaly_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We could superimpose (overlay) the anomaly map on top of the original image to get a heat map. Anomalib has a built-in function to achieve this. Let's try it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "heat_map = superimpose_anomaly_map(anomaly_map=anomaly_map, image=image, normalize=True)\n",
    "plt.imshow(heat_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "`predictions` also contains prediction scores and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pred_score = predictions[\"pred_scores\"][0]\n",
    "pred_labels = predictions[\"pred_labels\"][0]\n",
    "print(pred_score, pred_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The last part of the predictions is the mask that is predicted by the model. This is a boolean mask containing True/False for the abnormal/normal pixels, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pred_masks = predictions[\"pred_masks\"][0].squeeze().cpu().numpy()\n",
    "plt.imshow(pred_masks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "That wraps it! In this notebook, we show how we could train, test and finally infer a FastFlow model using Anomalib API."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anomalib",
   "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.8.13 (default, Oct 21 2022, 23:50:54) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f26beec5b578f06009232863ae217b956681fd13da2e828fa5a0ecf8cf2ccd29"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}