recognition.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8dec0b18-a8d5-432d-9f42-925ecb3f0e65",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cbmi/miniconda3/envs/wyx/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import random\n",
    "import argparse\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "from glob import glob\n",
    "from multiprocessing import Pool\n",
    "from torch.optim import optimizer\n",
    "from collections import OrderedDict\n",
    "from torch.utils.data import DataLoader\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "039a14e3-0d26-482d-ab90-1063cf9a5f49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.common_utils.dataset_building import split_multi_individuals_datasets\n",
    "from src.common_utils.prints import get_checkpoint_timestamp, print_log_message\n",
    "from src.preprocessing.training.utils import extract_annos\n",
    "from src.recognition.training.configs import e1 as cfg\n",
    "from src.recognition.training.extract_training_data import extract_preprocessing_json, select_ids\n",
    "from src.recognition.inference.feature_maker.load_features import make_one_volume_neuronal_features\n",
    "from src.common_utils.prints import print_info_message\n",
    "from src.common_utils.metric.rec import top_1_accuracy_score, top_k_accuracy, top_1_accuracy_score_torch\n",
    "from src.recognition.training.extract_training_data import neurons2data\n",
    "from src.recognition.training.exps.e1 import *\n",
    "from src.recognition.inference.dataset import RecFeatureDataset\n",
    "from src.recognition.inference.network import RecFuseNetworkLinear, RecMarginalCosLossNetwork"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "60656ec6-f1c9-4c67-b580-3c6488859ada",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f344c09a0d0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser = argparse.ArgumentParser(description = '')\n",
    "parser.add_argument('--name-reg', type = str, default = r\"[iI]ma?ge?_?[sS]t(?:ac)?k_?\\d+_dk?\\d+.*[wW]\\d+_?Dt\\d{6}\")\n",
    "parser.add_argument('--load-preprocess-result-root', type = str, default = \"data/dataset/proofreading\")\n",
    "parser.add_argument('--data-root', type = str, default = \"\")\n",
    "parser.add_argument('--random-seed', type = int, default = 520)\n",
    "parser.add_argument('--checkpoint-timestamp', type = str, default = f\"{get_checkpoint_timestamp()}_{np.random.randint(100000)}\")\n",
    "parser.add_argument('--label-root', type = str, default = \"data/dataset/label\")\n",
    "# neural points setting\n",
    "parser.add_argument('--rec-tflag', default = 130, type = int)\n",
    "parser.add_argument('--rec-fea-mode', default = 0, type = int)\n",
    "parser.add_argument('--rec-others-class', default = 1, type = int)\n",
    "parser.add_argument('--rec-xoy-unit', type = float, default = 0.3, help = \"um/pixel\")\n",
    "parser.add_argument('--rec-z-unit', type = float, default = 1.5, help = \"um/pixel\")\n",
    "parser.add_argument('--rec-worm-diagonal-line', type = float, default = 400.0)\n",
    "# knn feature\n",
    "parser.add_argument('--rec-knn-k', type = int, default = 25)\n",
    "# neural density feature\n",
    "parser.add_argument('--rec-des-len', type = int, default = 20)\n",
    "# neuron recognition (train)\n",
    "parser.add_argument('--rec-fp16', action = \"store_true\")\n",
    "parser.add_argument('--rec-epoch', default = 300, type = int)\n",
    "parser.add_argument('--rec-num-workers', default = 8, type = int)\n",
    "parser.add_argument('--rec-batch-size', default = 256, type = int)\n",
    "parser.add_argument('--rec-model-load-path', type = str, default = \"\")\n",
    "parser.add_argument('--rec-model-save-path', type = str, default = \"models/supp/e1\")\n",
    "parser.add_argument('--rec-shuffle', default = 1, type = int)\n",
    "\n",
    "# embedding method\n",
    "parser.add_argument('--rec-channel-base', type = int, default = 32)\n",
    "parser.add_argument('--rec-group-base', type = int, default = 4)\n",
    "parser.add_argument('--rec-len-embedding', type = int, default = 56)\n",
    "parser.add_argument('--rec-hypersphere-radius', type = int, default = 32)\n",
    "parser.add_argument('--rec-loss-coefficients', type = float, nargs = \"+\", default = [1.05, 0.0, 0.05])\n",
    "# tensorboard\n",
    "parser.add_argument('--rec-tensorboard-root', type = str, default = \"tb_log/supp/e1\")\n",
    "\n",
    "args = parser.parse_args(\"\")\n",
    "args.rec_z_scale = args.rec_z_unit / args.rec_xoy_unit\n",
    "\n",
    "random.seed(args.random_seed)\n",
    "np.random.seed(args.random_seed)\n",
    "torch.manual_seed(args.random_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fc4da9c4-3b65-4d4b-a14f-b12197d0ca3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- data preparation --------------------\n",
    "# vols_ccords: {vol_name: [[xmin, ymin], mass_of_center, anterior_y, posterior_y, ventral_x, dorsal_x]}\n",
    "vols_xymin, vols_ccords = extract_preprocessing_json(os.path.join(args.load_preprocess_result_root, \"*/*.json\"))\n",
    "\n",
    "labels = [{k: v for file_name, idxes in idv_label for k, v in extract_annos(os.path.join(args.label_root, file_name), idxes, args.name_reg).items()}\n",
    "          for idv_name, idv_label in cfg.dataset['animals']['label'].items()]\n",
    "labels = [{k: {i: [[p[1] - vols_xymin[k][0], p[2] - vols_xymin[k][1], p[3] - vols_xymin[k][0], p[4] - vols_xymin[k][1], p[0]] for p in pp] for i, pp in vol.items()}\n",
    "           for k, vol in idv_labels.items()} for idv_labels in labels]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f015710f-6181-49ed-a66d-d6dbf5fc469d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========  within dataset  ====================\n",
    "t_idx = args.rec_tflag\n",
    "val_idx = t_idx + 20\n",
    "processing_ids = select_ids(labels[0], 225)[0]\n",
    "within_labels, dataset_names = split_multi_individuals_datasets(labels = labels[:1], indexes = [[(0, t_idx), (t_idx, val_idx), (val_idx, len(labels[0]))]], shuffle_type = args.rec_shuffle)\n",
    "# ----- samples making --------------------\n",
    "within_engineering_feature = make_fea_multiprocess(within_labels, vols_ccords, args, mode = args.rec_fea_mode)\n",
    "trains, vals, tests, num_ids, id_map, processing_ids = neurons2data(within_engineering_feature.copy(), dataset_names = dataset_names,\n",
    "                                                                    include_others_class = args.rec_others_class, given_ids = sorted(processing_ids), verbose = False)\n",
    "# ----- dataloader --------------------\n",
    "train_dataset = RecFeatureDataset(Xs = trains[0], ys = trains[1], names = trains[2], is_train = True, is_fp16 = args.rec_fp16)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size = args.rec_batch_size, drop_last = True, shuffle = True, pin_memory = False, num_workers = 1)\n",
    "val_dataset = RecFeatureDataset(Xs = vals[0], ys = vals[1], names = vals[2], is_train = False, is_fp16 = args.rec_fp16)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size = args.rec_batch_size, drop_last = False, shuffle = True, pin_memory = False, num_workers = 1)\n",
    "within_test_dataset = RecFeatureDataset(tests[0], tests[1], tests[2], is_train = False, is_fp16 = args.rec_fp16)\n",
    "within_test_dataloader = DataLoader(within_test_dataset, batch_size = args.rec_batch_size, drop_last = False, shuffle = True, pin_memory = False, num_workers = 1)\n",
    "\n",
    "# ==========  across dataset  ====================\n",
    "across_test_infos = [make_across_testset(make_fea_multiprocess(idv_label, vols_ccords, args)) for idv_label in labels[1:]]\n",
    "across_test_datasets = [RecFeatureDataset(info[0], info[1], info[2], is_train = False, is_fp16 = args.rec_fp16) for info in across_test_infos]\n",
    "across_test_vol_names = [info[3] for info in across_test_infos]\n",
    "across_test_dataloaders = [DataLoader(testset, batch_size = args.rec_batch_size) for testset in across_test_datasets]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0911b6cc-43d4-4381-828b-71254e0c852a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- network --------------------\n",
    "model = RecFuseNetworkLinear(input_dim = (args.rec_knn_k * 3, args.rec_des_len * 4), output_dim = args.rec_len_embedding, num_ids = num_ids,\n",
    "                             channel_base = args.rec_channel_base, group_base = args.rec_group_base,\n",
    "                             dropout_ratio = 0.2, activation_method = \"celu\").cuda()\n",
    "model = model.half() if args.rec_fp16 else model\n",
    "if os.path.isfile(args.rec_model_load_path):\n",
    "    model.load_state_dict(torch.load(args.rec_model_load_path, map_location = 'cuda:0')['network'])\n",
    "    \n",
    "criterion = RecMarginalCosLossNetwork(len_embedding = args.rec_len_embedding, coefficients = args.rec_loss_coefficients, hypersphere_radius = args.rec_hypersphere_radius).cuda()\n",
    "criterion = criterion.half() if args.rec_fp16 else criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d564f6ab-c328-41c6-9dde-c132f17f7068",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-07-29 09:25:50 - \u001b[32m\u001b[1mINFO   \u001b[0m - ------------------------------------------- >>>>>>>>>>>>>>>>>>>>>>>>>>>> Training procedure starts! \n",
      "\n",
      "2022-07-29 09:35:29 - \u001b[32m\u001b[1mINFO   \u001b[0m - ------------------------------------------- >>>>>>>>>>>>>>>>>>>>>>>>>>>> Training procedure finished! \n",
      "The best val epoch of 2022_07_29_09_25_26_8335 model is 260 \t Top-1 accuracy: 87.68 % \t \n",
      "2022-07-29 09:35:30 - \u001b[32m\u001b[1mINFO   \u001b[0m - Within testset: accuracy: 95.48 \t num_hits: 150.67\n",
      "2022-07-29 09:35:31 - \u001b[32m\u001b[1mINFO   \u001b[0m - Across testset: 2 animals \t accuracy: 66.70 \t num_gts: 139.19 \t num_hits: 92.84\n"
     ]
    }
   ],
   "source": [
    "train_val_procedure(model = model,\n",
    "                    criterion = criterion,\n",
    "                    train_dataloader = train_dataloader,\n",
    "                    val_dataloader = val_dataloader,\n",
    "                    within_test_dataloader = within_test_dataloader,\n",
    "                    across_test_dataloaders = across_test_dataloaders,\n",
    "                    across_test_vol_names = across_test_vol_names,\n",
    "                    cfg = cfg,\n",
    "                    test_vols = tests[3],\n",
    "                    batch_size = args.rec_batch_size,\n",
    "                    optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3),\n",
    "                    is_fp16 = args.rec_fp16,\n",
    "                    num_ids = num_ids,\n",
    "                    processing_ids = processing_ids,\n",
    "                    hr = args.rec_hypersphere_radius,\n",
    "                    num_epochs = args.rec_epoch,\n",
    "                    checkpoint_timestamp = args.checkpoint_timestamp,\n",
    "                    tb_writer = SummaryWriter(os.path.join(args.rec_tensorboard_root, args.checkpoint_timestamp)),\n",
    "                    model_save_path = os.path.join(args.rec_model_save_path, f\"rec_{args.checkpoint_timestamp}.ckpt\"),\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "778c48c3-c332-4a6e-bb11-bb00fe1d257c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-07-29 09:35:36 - \u001b[32m\u001b[1mINFO   \u001b[0m - NeRVE: accuracy:  50.08 \t num_gts: 61.41 \t num_hits: 30.76\n",
      "2022-07-29 09:35:37 - \u001b[32m\u001b[1mINFO   \u001b[0m - NeuroPAL Yu: accuracy: 27.79 ± 1.11 \t num_gts: 38.81 ± 2.87 \t num_hits: 10.82 ± 1.32\n",
      "2022-07-29 09:35:38 - \u001b[32m\u001b[1mINFO   \u001b[0m - NeuroPAL Chaudhary: accuracy: 42.97 ± 2.35 \t num_gts: 49.62 ± 0.71 \t num_hits: 21.31 ± 0.90\n"
     ]
    }
   ],
   "source": [
    "# ==========  evaluate benchmarks  ====================\n",
    "from src.benchmarks.test2 import evaluate_benchmark\n",
    "from src.benchmarks.datasets.CeNDeR import Dataset_CeNDeR\n",
    "\n",
    "model.eval()\n",
    "test_batch_size = 32  # keep the same with fDNC\n",
    "\n",
    "# benchmark leifer 2017\n",
    "dataset = Dataset_CeNDeR(glob(os.path.join(args.data_root, \"data/benchmarks/supp/e1/test_tracking\", \"*.npy\")), is_fp16 = args.rec_fp16, num_pool = args.rec_num_workers)\n",
    "dataloader = DataLoader(dataset, args.rec_batch_size)\n",
    "accs = evaluate_benchmark(dataloader, model, refer_idx = 16, verbose = False, test_batch_size = test_batch_size)\n",
    "print_info_message(f\"NeRVE: accuracy:  {accs[0] * 100:.2f} \\t num_gts: {accs[2]:.2f} \\t num_hits: {(accs[4]):.2f}\")\n",
    "\n",
    "# benchmark NeuroPAL\n",
    "dataset = Dataset_CeNDeR(glob(os.path.join(args.data_root, \"data/benchmarks/supp/e1/test_neuropal_our\", \"*.npy\")), is_fp16 = args.rec_fp16, num_pool = args.rec_num_workers)\n",
    "dataloader = DataLoader(dataset, args.rec_batch_size)\n",
    "accs = np.array([evaluate_benchmark(dataloader, model, refer_idx = ref_idx, test_batch_size = test_batch_size) for ref_idx in range(dataset.num_vols)])\n",
    "print_info_message(f\"NeuroPAL Yu: accuracy: {np.mean(accs[:, 0] * 100):.2f} ± {np.std(accs[:, 0] * 100):.2f} \\t num_gts: {np.mean(accs[:, 2]):.2f} ± {np.std(accs[:, 2]):.2f} \\t \"\n",
    "                   f\"num_hits: {np.mean(accs[:, 4]):.2f} ± {np.std(accs[:, 4]):.2f}\")\n",
    "\n",
    "# benchmark NeuroPAL Chaudhary\n",
    "dataset = Dataset_CeNDeR(glob(os.path.join(args.data_root, \"data/benchmarks/supp/e1/test_neuropal_Chaudhary\", \"*.npy\")), is_fp16 = args.rec_fp16, num_pool = args.rec_num_workers)\n",
    "dataloader = DataLoader(dataset, args.rec_batch_size)\n",
    "accs = np.array([evaluate_benchmark(dataloader, model, refer_idx = ref_idx, test_batch_size = test_batch_size) for ref_idx in range(dataset.num_vols)])\n",
    "print_info_message(f\"NeuroPAL Chaudhary: accuracy: {np.mean(accs[:, 0] * 100):.2f} ± {np.std(accs[:, 0] * 100):.2f} \\t num_gts: {np.mean(accs[:, 2]):.2f} ± {np.std(accs[:, 2]):.2f} \\t \"\n",
    "                   f\"num_hits: {np.mean(accs[:, 4]):.2f} ± {np.std(accs[:, 4]):.2f}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wyx",
   "language": "python",
   "name": "wyx"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}