{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Skeletonization: supervised, inference\n",
    "======================================\n",
    "\n",
    "In this notebook we use the supervised module to extract length and head width using a model trained on manually annotated data. We will use the script `skeletons/main_supervised_skeletons_inference.py` to extract skeletons form the clips. \n",
    "\n",
    "We first import the necessary libraries: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import cv2\n",
    "\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "from skimage.morphology import thin\n",
    "from torchvision import transforms\n",
    "from tqdm import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import yaml\n",
    "\n",
    "from mzbsuite.skeletons.mzb_skeletons_pilmodel import MZBModel_skels\n",
    "from mzbsuite.skeletons.mzb_skeletons_helpers import paint_image_tensor, Denormalize\n",
    "from mzbsuite.utils import cfg_to_arguments, find_checkpoints\n",
    "\n",
    "# Set the thread layer used by MKL\n",
    "os.environ[\"MKL_THREADING_LAYER\"] = \"GNU\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to set up some running parameters for the script too: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'config_file': WindowsPath('D:/mzb-workflow/configs/mzb_example_config.yaml'), 'input_dir': WindowsPath('D:/mzb-workflow/data/bgb/derived/blobs'), 'input_type': 'external', 'input_model': WindowsPath('D:/mzb-workflow/models/mzb-skeleton-models/mit-b2-v1'), 'output_dir': WindowsPath('D:/mzb-workflow/results/bgb/skeletons/skeletons_supervised'), 'save_masks': WindowsPath('D:/mzb-workflow/data/bgb/skeletons/skeletons_supervised'), 'verbose': True}\n"
     ]
    }
   ],
   "source": [
    "ROOT_DIR = ROOT_DIR = Path(\"D:\\mzb-workflow\") #Path(\"/data/shared/mzb-workflow\")\n",
    "MODEL = \"mit-b2-v1\"\n",
    "\n",
    "arguments = {\n",
    "    \"config_file\": ROOT_DIR / \"configs/mzb_example_config.yaml\",\n",
    "    \"input_dir\": ROOT_DIR / \"data/bgb/derived/blobs\",\n",
    "    \"input_type\": \"external\", \n",
    "    \"input_model\": ROOT_DIR / f\"models/mzb-skeleton-models/{MODEL}\", \n",
    "    \"output_dir\": ROOT_DIR / \"results/bgb/skeletons/skeletons_supervised\",\n",
    "    \"save_masks\": ROOT_DIR / \"data/bgb/skeletons/skeletons_supervised\", \n",
    "    \"verbose\": True,\n",
    "}\n",
    "    \n",
    "with open(str(arguments[\"config_file\"]), \"r\") as f:\n",
    "    cfg = yaml.load(f, Loader=yaml.FullLoader)\n",
    "\n",
    "# cfg[\"trcl_gpu_ids\"] = None\n",
    "print(arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert to a dictionary for the scripts to parse. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'glob_random_seed': 222, 'glob_root_folder': '/home/jovyan/work/mzb-workflow/', 'glob_blobs_folder': '/home/jovyan/work/mzb-workflow/data/derived/blobs/', 'glob_local_format': 'pdf', 'model_logger': 'wandb', 'impa_image_format': 'jpg', 'impa_clip_areas': [2700, 4700, -1, -1], 'impa_area_threshold': 5000, 'impa_gaussian_blur': [21, 21], 'impa_gaussian_blur_passes': 3, 'impa_adaptive_threshold_block_size': 351, 'impa_mask_postprocess_kernel': [11, 11], 'impa_mask_postprocess_passes': 5, 'impa_bounding_box_buffer': 200, 'impa_save_clips_plus_features': True, 'lset_class_cut': 'order', 'lset_val_size': 0.1, 'trcl_learning_rate': 0.0001, 'trcl_batch_size': 8, 'trcl_weight_decay': 0, 'trcl_step_size_decay': 5, 'trcl_number_epochs': 75, 'trcl_save_topk': 1, 'trcl_num_classes': 8, 'trcl_model_pretrarch': 'convnext-small', 'trcl_num_workers': 16, 'trcl_wandb_project_name': 'mzb-classifiers', 'trcl_logger': 'wandb', 'trsk_learning_rate': 0.001, 'trsk_batch_size': 32, 'trsk_weight_decay': 0, 'trsk_step_size_decay': 25, 'trsk_number_epochs': 400, 'trsk_save_topk': 1, 'trsk_num_classes': 2, 'trsk_model_pretrarch': 'mit_b2', 'trsk_num_workers': 16, 'trsk_wandb_project_name': 'mzb-skeletons', 'trsk_logger': 'wandb', 'infe_model_ckpt': 'last', 'infe_num_classes': 8, 'infe_image_glob': '*_rgb.jpg', 'skel_class_exclude': 'errors', 'skel_conv_rate': 131.6625, 'skel_label_thickness': 3, 'skel_label_buffer_on_preds': 25, 'skel_label_clip_with_mask': False}\n"
     ]
    }
   ],
   "source": [
    "# Transforms configurations dicts to argparse arguments\n",
    "args = cfg_to_arguments(arguments)\n",
    "cfg = cfg_to_arguments(cfg)\n",
    "print(str(cfg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can load the code necessary to run the inference from the dedicated script, and call it with the arguments specified above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;31mSignature:\u001b[0m \u001b[0minference_skeleton\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcfg\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mDocstring:\u001b[0m\n",
      "Function to run inference of skeletons (body, head) on macrozoobenthos images clips, using a trained model.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "args : argparse.Namespace\n",
      "    Namespace containing the arguments passed to the script. Notably:\n",
      "\n",
      "        - input_dir: path to the directory containing the images to be classified\n",
      "        - input_type: type of input data, either \"val\" or \"external\"\n",
      "        - input_model: path to the directory containing the model to be used for inference\n",
      "        - output_dir: path to the directory where the results will be saved\n",
      "        - save_masks: path to the directory where the masks will be saved\n",
      "        - config_file: path to the config file with train / inference parameters\n",
      "\n",
      "cfg : dict\n",
      "    Dictionary containing the configuration parameters.\n",
      "\n",
      "Returns\n",
      "-------\n",
      "None. Saves the results in the specified folder.\n",
      "\u001b[1;31mFile:\u001b[0m      d:\\mzb-workflow\\scripts\\skeletons\\main_supervised_skeleton_inference.py\n",
      "\u001b[1;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "# from classification.main_classification_finetune import main as finetune_classifier\n",
    "from scripts.skeletons.main_supervised_skeleton_inference import main as inference_skeleton\n",
    "?inference_skeleton"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can call the function and run the inference on the images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 13/13 [00:05<00:00,  2.35it/s]\n",
      "Neural network predictions done, refining and saving skeletons...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [02:49<00:00,  1.70s/it]\n"
     ]
    }
   ],
   "source": [
    "inference_skeleton(args, cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This produces a `.csv` file with the predictions for body length and head width (saved in `output_dir`), as well as the predicted skeletons for body length and head for each clip (saved in `save_masks`). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualise the results and assess model accuracy against a manually annotated validation set, if available. First we need to provide some additional arguments: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'config_file': WindowsPath('D:/mzb-workflow/configs/mzb_example_config.yaml'), 'input_dir': WindowsPath('D:/mzb-workflow/data/mzb_example_data/derived/blobs'), 'input_type': 'external', 'input_model': WindowsPath('D:/mzb-workflow/models/mzb-skeleton-models/mit-b2-v1'), 'output_dir': WindowsPath('D:/mzb-workflow/results/mzb_example_data/skeletons/skeletons_supervised'), 'save_masks': WindowsPath('D:/mzb-workflow/data/bgb/skeletons/skeletons_supervised'), 'verbose': True, 'manual_annotations': WindowsPath('D:/mzb-workflow/data/mzb_example/skeletons/supervised_skeletons/manual_anns/manual_annotations_summary.csv'), 'model_annotations': WindowsPath('D:/mzb-workflow/results/mzb_example_data/skeletons/skeletons_supervised/size_skel_supervised_model.csv')}\n"
     ]
    }
   ],
   "source": [
    "arguments['input_dir'] = ROOT_DIR / \"data/mzb_example_data/derived/blobs\"\n",
    "arguments['manual_annotations'] = ROOT_DIR / \"data\\mzb_example\\skeletons\\supervised_skeletons\\manual_anns\\manual_annotations_summary.csv\"\n",
    "arguments['model_annotations'] = ROOT_DIR / \"results\\mzb_example_data\\skeletons\\skeletons_supervised\\size_skel_supervised_model.csv\"\n",
    "arguments['output_dir'] = ROOT_DIR / \"results\\mzb_example_data\\skeletons\\skeletons_supervised\"\n",
    "print(arguments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'config_file': WindowsPath('D:/mzb-workflow/configs/mzb_example_config.yaml'), 'input_dir': WindowsPath('D:/mzb-workflow/data/mzb_example_data/derived/blobs'), 'input_type': 'external', 'input_model': WindowsPath('D:/mzb-workflow/models/mzb-skeleton-models/mit-b2-v1'), 'output_dir': WindowsPath('D:/mzb-workflow/results/mzb_example_data/skeletons/skeletons_supervised'), 'save_masks': WindowsPath('D:/mzb-workflow/data/bgb/skeletons/skeletons_supervised'), 'verbose': True, 'manual_annotations': WindowsPath('D:/mzb-workflow/data/mzb_example/skeletons/supervised_skeletons/manual_anns/manual_annotations_summary.csv'), 'model_annotations': WindowsPath('D:/mzb-workflow/results/mzb_example_data/skeletons/skeletons_supervised/size_skel_supervised_model.csv')}\n"
     ]
    }
   ],
   "source": [
    "args = cfg_to_arguments(arguments)\n",
    "print(str(args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;31mSignature:\u001b[0m \u001b[0massess_skeletons\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcfg\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mDocstring:\u001b[0m\n",
      "Main function to run an assessment of the length measurements.\n",
      "Computes the absolute error between manual annotations and model predictions, and reports plots grouped by species.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "args: argparse.Namespace\n",
      "    Arguments parsed from the command line. Specifically:\n",
      "    \n",
      "        - args.input_dir: path to the directory with the model predictions\n",
      "        - args.manual_annotations: path to the manual annotations\n",
      "        - args.model_annotations: path to the model predictions\n",
      "        - args.output_dir: path to the directory where plots accuracy report should be saved\n",
      "\n",
      "cfg: argparse.Namespace\n",
      "    configuration options.\n",
      "\n",
      "Returns\n",
      "-------\n",
      "None. Plots and metrics are saved in the results directory.\n",
      "\u001b[1;31mFile:\u001b[0m      d:\\mzb-workflow\\scripts\\skeletons\\main_supervised_skeleton_assessment.py\n",
      "\u001b[1;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "from scripts.skeletons.main_supervised_skeleton_assessment import main as assess_skeletons\n",
    "?assess_skeletons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assess_skeletons(args, cfg)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}