show.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "from pathlib import Path\n",
    "from matplotlib import ticker\n",
    "import matplotlib.patheffects as pe\n",
    "from matplotlib.offsetbox import OffsetImage, AnnotationBbox\n",
    "from multiprocessing import Pool\n",
    "from tqdm.auto import tqdm\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "import scipy\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "\n",
    "from skimage import exposure, transform\n",
    "from skimage.draw import line_aa\n",
    "from skimage.color import gray2rgb, rgb2gray\n",
    "from skimage.io import imread, imsave\n",
    "from skimage.transform import resize\n",
    "from skimage.util import montage\n",
    "\n",
    "from sklearn.metrics import jaccard_score\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "from datasets import PerineuronalNetsDataset, CellsDataset, PerineuronalNetsRankDataset\n",
    "\n",
    "from methods.points.match import match\n",
    "from methods.points.metrics import detection_and_counting\n",
    "from methods.points.utils import draw_points, draw_groundtruth_and_predictions\n",
    "\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VGG Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "common = dict(root='data/vgg-cells')\n",
    "\n",
    "detection_kws = {\n",
    "    'target_': 'detection',\n",
    "    'target_params': {\n",
    "        'side': 12\n",
    "    }\n",
    "}\n",
    "\n",
    "density_kws = {\n",
    "    'target_': 'density',\n",
    "    'target_params': {\n",
    "        'mode': 'reflect',\n",
    "        'k_size': 51,\n",
    "        'sigma': 5\n",
    "    }\n",
    "}\n",
    "\n",
    "segmentation_kws = {\n",
    "    'target_': 'segmentation',\n",
    "    'target_params': {\n",
    "        'radius': 5,\n",
    "        'radius_ignore': 6,\n",
    "        'v_bal': 0.1, \n",
    "        'sigma_bal': 3,\n",
    "        'sep_width': 1,     \n",
    "        'sigma_sep': 3,\n",
    "        'lambda_sep': 50\n",
    "    }\n",
    "}\n",
    "\n",
    "detection_dataset = CellsDataset(**detection_kws, **common)\n",
    "density_dataset = CellsDataset(**density_kws, **common)\n",
    "segmentation_dataset = CellsDataset(**segmentation_kws, **common)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_idx = 0\n",
    "\n",
    "sample = density_dataset[sample_idx][0][:, :, 0]\n",
    "detections = detection_dataset[sample_idx][0][1]\n",
    "density_map = density_dataset[sample_idx][0][:, :, 1]\n",
    "segmentation_map = segmentation_dataset[sample_idx][0][:, :, 1]\n",
    "weights_map = segmentation_dataset[sample_idx][0][:, :, 2]\n",
    "\n",
    "detections = np.clip(detections.astype(int), 0, sample.shape[0] - 1)\n",
    "\n",
    "sample = matplotlib.cm.jet(sample)\n",
    "density_map = gray2rgb(density_map / density_map.max())\n",
    "segmentation_map = gray2rgb(segmentation_map)\n",
    "\n",
    "sample_with_boxes = sample.copy()\n",
    "for y0, x0, y1, x1 in detections:\n",
    "    rect = ((y0, x0, y0, x1),\n",
    "            (y0, x1, y1, x1),\n",
    "            (y1, x1, y1, x0),\n",
    "            (y1, x0, y0, x0))\n",
    "    for r0, c0, r1, c1 in rect:\n",
    "        rr, cc, val = line_aa(r0, c0, r1, c1)\n",
    "        sample_with_boxes[rr, cc, 0] = val\n",
    "\n",
    "        \n",
    "fig, axes = plt.subplots(1, 4, figsize=(20, 20))\n",
    "axes = axes.flatten()\n",
    "axes[0].imshow(sample)\n",
    "axes[1].imshow(sample_with_boxes)\n",
    "axes[2].imshow(density_map)\n",
    "axes[3].imshow(segmentation_map)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.set_axis_off()\n",
    "\n",
    "plt.imsave('figures/vgg-sample.png', sample)\n",
    "plt.imsave('figures/vgg-boxes.png', sample_with_boxes)\n",
    "plt.imsave('figures/vgg-density.png', density_map)\n",
    "plt.imsave('figures/vgg-segmentation.png', segmentation_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncells = [len(detection_dataset[i][0][1]) for i in range(len(detection_dataset)) ]\n",
    "f'{sum(ncells)} cells  {np.mean(ncells):.2f}$\\pm${np.std(ncells):.2f} cells/image'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_runs(model_name, run, csv_file):\n",
    "    run = Path(run)\n",
    "    cfg = OmegaConf.load(run / '.hydra' / 'config.yaml')\n",
    "    num_samples = cfg['data']['validation']['num_samples'][0]\n",
    "    seed = cfg['data']['validation']['split_seed']\n",
    "\n",
    "    csv_path = run / 'test_predictions' / csv_file\n",
    "    if not csv_path.exists():\n",
    "        print(f'Skipping not found: {csv_path}')\n",
    "        return pd.DataFrame()\n",
    "    \n",
    "    data = pd.read_csv(csv_path, index_col=0)\n",
    "    data['model'] = model_name\n",
    "    data['num_samples'] = num_samples\n",
    "    data['split_seed'] = seed\n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runs = {\n",
    "    'S-UNet': Path('runs/experiment=vgg-cells/segmentation/').glob('unet_*'),\n",
    "    'FRCNN': Path('runs/experiment=vgg-cells/detection/').glob('fasterrcnn_*'),\n",
    "    'D-CSRNet': Path('runs/experiment=vgg-cells/density/').glob('csrnet_*')\n",
    "}\n",
    "\n",
    "metrics = pd.concat([collect_runs(k, r, 'all_metrics.csv.gz') for k, v in runs.items() for r in v], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_metrics = metrics.groupby(['model', 'num_samples', 'split_seed', 'thr'])['count/mae'].mean()\n",
    "best_configs = mean_metrics.groupby(['model', 'num_samples', 'split_seed']).idxmin()\n",
    "table = mean_metrics.loc[best_configs].groupby(['model', 'num_samples']).apply(lambda x: pd.Series({'mean': x.mean(), 'std': x.std()}))\n",
    "table = table.unstack(2).apply(lambda x: f'{x[\"mean\"]:.1f} $\\pm$ {x[\"std\"]:.1f}', axis=1).unstack(1)\n",
    "print(table.to_latex(escape=False))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MBM Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "common = dict(root='data/mbm-cells')\n",
    "\n",
    "detection_kws = {\n",
    "    'target_': 'detection',\n",
    "    'target_params': {\n",
    "        'side': 20\n",
    "    }\n",
    "}\n",
    "\n",
    "density_kws = {\n",
    "    'target_': 'density',\n",
    "    'target_params': {\n",
    "        'mode': 'reflect',\n",
    "        'k_size': 51,\n",
    "        'sigma': 10\n",
    "    }\n",
    "}\n",
    "\n",
    "segmentation_kws = {\n",
    "    'target_': 'segmentation',\n",
    "    'target_params': {\n",
    "        'radius': 12,\n",
    "        'radius_ignore': 15,\n",
    "        'v_bal': 0.1, \n",
    "        'sigma_bal': 5,\n",
    "        'sep_width': 1,     \n",
    "        'sigma_sep': 4,\n",
    "        'lambda_sep': 50\n",
    "    }\n",
    "}\n",
    "\n",
    "detection_dataset = CellsDataset(**detection_kws, **common)\n",
    "density_dataset = CellsDataset(**density_kws, **common)\n",
    "segmentation_dataset = CellsDataset(**segmentation_kws, **common)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_idx = 5\n",
    "\n",
    "# sample = density_dataset[sample_idx][0][:, :, 0]\n",
    "sample = imread('data/mbm-cells/BM_GRAZ_HE_0001_02_001_cell.png')\n",
    "detections = detection_dataset[sample_idx][0][1]\n",
    "density_map = density_dataset[sample_idx][0][:, :, 1]\n",
    "segmentation_map = segmentation_dataset[sample_idx][0][:, :, 1]\n",
    "weights_map = segmentation_dataset[sample_idx][0][:, :, 2]\n",
    "\n",
    "detections = np.clip(detections.astype(int), 0, sample.shape[0] - 1)\n",
    "sw_map = np.stack((segmentation_map, weights_map, np.zeros_like(weights_map)), axis=-1)\n",
    "\n",
    "a = np.stack((rgb2gray(sample), segmentation_map, np.zeros_like(segmentation_map)), axis=-1)\n",
    "\n",
    "density_map = gray2rgb(density_map / density_map.max())\n",
    "segmentation_map = gray2rgb(segmentation_map)\n",
    "\n",
    "sample_with_boxes = sample.copy()\n",
    "for y0, x0, y1, x1 in detections:\n",
    "    rect = ((y0, x0, y0, x1),\n",
    "            (y0, x1, y1, x1),\n",
    "            (y1, x1, y1, x0),\n",
    "            (y1, x0, y0, x0))\n",
    "    for r0, c0, r1, c1 in rect:\n",
    "        rr, cc, val = line_aa(r0, c0, r1, c1)\n",
    "        sample_with_boxes[rr, cc, 0] = val\n",
    "\n",
    "        \n",
    "fig, axes = plt.subplots(1, 4, figsize=(20, 20))\n",
    "axes = axes.flatten()\n",
    "axes[0].imshow(sample)\n",
    "axes[1].imshow(sample_with_boxes)\n",
    "axes[2].imshow(density_map)\n",
    "axes[3].imshow(segmentation_map)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.set_axis_off()\n",
    "\n",
    "plt.subplots_adjust(wspace=0)\n",
    "\n",
    "plt.imsave('figures/mbm-sample.png', sample)\n",
    "plt.imsave('figures/mbm-boxes.png', sample_with_boxes)\n",
    "plt.imsave('figures/mbm-density.png', density_map)\n",
    "plt.imsave('figures/mbm-segmentation.png', segmentation_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runs = {\n",
    "    'S-UNet': Path('runs/experiment=mbm-cells/segmentation/').glob('unet_*'),\n",
    "    'FRCNN': Path('runs/experiment=mbm-cells/detection/').glob('fasterrcnn_*'),\n",
    "    'D-CSRNet': Path('runs/experiment=mbm-cells/density/').glob('csrnet_*')\n",
    "}\n",
    "\n",
    "metrics = pd.concat([collect_runs(k, r, 'all_metrics.csv.gz') for k, v in runs.items() for r in v], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_metrics = metrics.groupby(['model', 'num_samples', 'split_seed', 'thr'])['count/mae'].mean()\n",
    "best_configs = mean_metrics.groupby(['model', 'num_samples', 'split_seed']).idxmin()\n",
    "table = mean_metrics.loc[best_configs].groupby(['model', 'num_samples']).apply(lambda x: pd.Series({'mean': x.mean(), 'std': x.std()}))\n",
    "table = table.unstack(2).apply(lambda x: f'{x[\"mean\"]:.1f} $\\pm$ {x[\"std\"]:.1f}', axis=1).unstack(1)\n",
    "print(table.to_latex(escape=False))\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PNN Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "common = dict(split='train-half1', random_offset=0, patch_size=640)\n",
    "detection_kws = {\n",
    "    'target_': 'detection',\n",
    "    'target_params': {\n",
    "        'side': 45\n",
    "    }\n",
    "}\n",
    "\n",
    "density_kws = {\n",
    "    'target_': 'density',\n",
    "    'target_params': {\n",
    "        'mode': 'reflect',\n",
    "        'k_size': 151,\n",
    "        'sigma': 15\n",
    "    }\n",
    "}\n",
    "\n",
    "segmentation_kws = {\n",
    "    'target_': 'segmentation',\n",
    "    'target_params': {\n",
    "        'radius': 20,\n",
    "        'radius_ignore': 25,\n",
    "        'v_bal': 0.1, \n",
    "        'sigma_bal': 10,\n",
    "        'sep_width': 1,     \n",
    "        'sigma_sep': 6,\n",
    "        'lambda_sep': 50\n",
    "    }\n",
    "}\n",
    "\n",
    "detection_dataset = PerineuronalNetsDataset(**detection_kws, **common)\n",
    "density_dataset = PerineuronalNetsDataset(**density_kws, **common)\n",
    "segmentation_dataset = PerineuronalNetsDataset(**segmentation_kws, **common)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_samples = np.random.randint(0, len(detection_dataset), 10000)\n",
    "\n",
    "cells_per_sample = [len(detection_dataset[x][0][1]) for x in tqdm(random_samples)]\n",
    "best_samples = np.argsort(cells_per_sample)[::-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_idx = random_samples[best_samples[31]]\n",
    "\n",
    "sample = density_dataset[sample_idx][0][:, :, 0]\n",
    "detections = detection_dataset[sample_idx][0][1]\n",
    "density_map = density_dataset[sample_idx][0][:, :, 1]\n",
    "segmentation_map = segmentation_dataset[sample_idx][0][:, :, 1]\n",
    "weights_map = segmentation_dataset[sample_idx][0][:, :, 2]\n",
    "\n",
    "\n",
    "detections = np.clip(detections.astype(int), 0, sample.shape[0] - 1)\n",
    "sw_map = np.stack((segmentation_map, weights_map, np.zeros_like(sample)), axis=-1)\n",
    "\n",
    "sample = matplotlib.cm.viridis(sample)\n",
    "density_map = gray2rgb(density_map / density_map.max())\n",
    "segmentation_map = gray2rgb(segmentation_map)\n",
    "\n",
    "sample_with_boxes = sample.copy()\n",
    "for y0, x0, y1, x1 in detections:\n",
    "    rect = ((y0, x0, y0, x1),\n",
    "            (y0, x1, y1, x1),\n",
    "            (y1, x1, y1, x0),\n",
    "            (y1, x0, y0, x0))\n",
    "    for r0, c0, r1, c1 in rect:\n",
    "        rr, cc, val = line_aa(r0, c0, r1, c1)\n",
    "        sample_with_boxes[rr, cc, 0] = val\n",
    "\n",
    "        \n",
    "fig, axes = plt.subplots(1, 4, figsize=(20, 20))\n",
    "axes = axes.flatten()\n",
    "axes[0].imshow(sample)\n",
    "axes[1].imshow(sample_with_boxes)\n",
    "axes[2].imshow(density_map)\n",
    "axes[3].imshow(segmentation_map)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.set_axis_off()\n",
    "\n",
    "plt.imsave('figures/pnn-sample.png', sample)\n",
    "plt.imsave('figures/pnn-boxes.png', sample_with_boxes)\n",
    "plt.imsave('figures/pnn-density.png', density_map)\n",
    "plt.imsave('figures/pnn-segmentation.png', segmentation_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Groundtruth Properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt = pd.read_csv('data/perineuronal-nets/test/annotations.csv')\n",
    "gt['agreement'] = gt.loc[:, 'AV':'VT'].sum(axis=1)\n",
    "\n",
    "gt.groupby('imgName').X.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distribution of Agreement in the (Multi-Rater) Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "sns.set_theme(context='notebook', style='ticks')\n",
    "data = gt.agreement.value_counts().sort_index()\n",
    "_, _, autotexts = plt.pie(data.values, labels=data.index,\n",
    "                          autopct='{:.2g}%'.format, pctdistance=0.75,\n",
    "                          # colors=sns.color_palette('rocket', 7)\n",
    "                          colors=sns.color_palette('rocket', as_cmap=True)(np.linspace(0, 1, 8)[1:])\n",
    "                         )\n",
    "plt.ylabel('agreement')\n",
    "\n",
    "plt.setp(autotexts, size=12)\n",
    "for t in autotexts[:4]:\n",
    "    t.set_color('white')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pnn_cells = PerineuronalNetsRankDataset(mode='patches')\n",
    "\n",
    "x = enumerate(pnn_cells.annot.agreement.values)\n",
    "x = sorted(x, key=lambda x: x[1])\n",
    "x = itertools.groupby(x, key=lambda x: x[1])\n",
    "\n",
    "means = []\n",
    "for agreement, group in x:\n",
    "    samples = [i for i, _ in group]\n",
    "    images = [pnn_cells[i][0].astype(np.float32) / 255. for i in samples]\n",
    "    mean = np.mean(images, axis=0)\n",
    "    means.append(mean)\n",
    "    \n",
    "    sorted_samples = np.sum((images * mean), axis=(1,2)).argsort()\n",
    "    sorted_samples = np.array(samples)[sorted_samples][::-1]\n",
    "    \n",
    "    print(f'{agreement}:', sorted_samples[:12].tolist(), ',')\n",
    "\n",
    "pnn_means = np.stack(means)\n",
    "pnn_means = (pnn_means - np.min(pnn_means)) / (np.max(pnn_means) - np.min(pnn_means))\n",
    "pnn_means = matplotlib.cm.viridis(pnn_means)[:,:,:,:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = sorted(gt.agreement.values) + [100]\n",
    "a = np.array(a).reshape((112, -1))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "ax = sns.heatmap(a, vmin=0, vmax=7, square=True,\n",
    "            linewidths=0,\n",
    "            antialiased=True,\n",
    "            rasterized=True,\n",
    "            cbar=False, cbar_kws={\"orientation\": \"horizontal\", 'pad': 0.05, 'ticks': range(0, 8), 'drawedges': True})\n",
    "plt.xticks([])\n",
    "plt.ylabel('Rater\\'s Agreement')\n",
    "\n",
    "def find_pos(a):\n",
    "    pos = np.unique(a, return_counts=True)[1].cumsum()\n",
    "    pos = np.insert(pos, 0, 1)\n",
    "    pos = (pos[:-1] + pos[1:]) / 2\n",
    "    return pos\n",
    "\n",
    "l_pos = find_pos(a[:, 0])\n",
    "r_pos = find_pos(a[:, 1])\n",
    "\n",
    "_ = plt.yticks(l_pos, range(1, 8), rotation=0)\n",
    "\n",
    "_, x_limit = plt.xlim()\n",
    "y_limit, _ = plt.ylim()\n",
    "\n",
    "def shuf(l):\n",
    "    l = l[:]\n",
    "    random.shuffle(l)\n",
    "    return l\n",
    "\n",
    "sample_indices = {    \n",
    "    1: [1565, 2300, 1913, 311, 1799, 763, 72],\n",
    "    2: [2309, 386, 983, 56, 286, 951, 1774],\n",
    "    3: [198, 1874, 392, 872, 78, 390, 1103],\n",
    "    4: [777, 219, 1944, 1066, 217, 1115, 96],\n",
    "    5: [220, 2174, 945, 389, 385, 1633, 593],\n",
    "    6: [2218, 2058, 1436, 2212, 2034, 1433, 207],\n",
    "    7: [453, 7, 6, 4, 12, 644, 20]\n",
    "}\n",
    "\n",
    "pct = gt.agreement.value_counts()\n",
    "pct = pct / pct.sum()\n",
    "pct = pct.sort_index().values\n",
    "\n",
    "cell_x = 1.51\n",
    "nr = 1\n",
    "for i, (ry, ly) in enumerate(zip(r_pos, l_pos), start=1):\n",
    "    cell_y = 1 - (i / len(sample_indices))\n",
    "\n",
    "    # percentage\n",
    "    ax.annotate(f'{pct[i - 1]:.0%}', (0.5, 1 - (ry+ly) / (2*y_limit)), xycoords='axes fraction',\n",
    "                color='white' if i < 4 else 'black', ha='center', va='center')\n",
    "    \n",
    "     # connector\n",
    "    ax.annotate('', xy=(1, 1 - ry / y_limit), xycoords='axes fraction',\n",
    "            xytext=(cell_x, cell_y + 0.06), textcoords='axes fraction',\n",
    "            arrowprops=dict(arrowstyle='-', color='0.2', connectionstyle='arc,angleA=0,angleB=0,armA=-7,armB=7,rad=0'))\n",
    "    \n",
    "    # mean image\n",
    "    imagebox = OffsetImage(pnn_means[i - 1], zoom=0.5, origin='upper')\n",
    "    ab = AnnotationBbox(imagebox, (cell_x, cell_y), xycoords='axes fraction', frameon=False, box_alignment=(0,0))\n",
    "    ax.add_artist(ab)\n",
    "    \n",
    "    # samples\n",
    "    cell_images = [pnn_cells[j][0] for j in sample_indices[i]]\n",
    "    cell_images = [matplotlib.cm.viridis(c) for c in cell_images]\n",
    "    cell_images = np.stack(cell_images)[:,:,:,:3]\n",
    "    image = montage(cell_images, grid_shape=(nr, len(cell_images) / nr), padding_width=5, fill=(1, 1, 1), multichannel=True)\n",
    "    image = image[5:-5, ...]\n",
    "    \n",
    "    imagebox = OffsetImage(image, zoom=0.5, origin='upper')\n",
    "    ab = AnnotationBbox(imagebox, (cell_x + 0.8, cell_y), xycoords='axes fraction', frameon=False, box_alignment=(0,0))\n",
    "    ax.add_artist(ab)\n",
    "    \n",
    "ax.annotate('mean', (cell_x, 1), xycoords='axes fraction')\n",
    "ax.annotate('samples', (cell_x + 2, 1), xycoords='axes fraction')\n",
    "plt.savefig('figures/pnn-mr-breakdown.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_id = gt.imgName.unique()[2]\n",
    "img = plt.imread('data/perineuronal-nets/test/fullFrames/' + sample_id)\n",
    "p2, p98 = np.percentile(img, (0.1, 99.9))\n",
    "img = exposure.rescale_intensity(img, in_range=(p2, p98))\n",
    "img = transform.resize(img, (1024, 1024))\n",
    "img = matplotlib.cm.viridis(img)[:,:,:3]\n",
    "\n",
    "scale_f = img.shape[0] / 2000\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.imshow(img)\n",
    "ax = plt.gca()\n",
    "ax.set_axis_off()\n",
    "\n",
    "colors = sns.color_palette('rocket', as_cmap=True)(np.linspace(0, 1, 8)[1:])\n",
    "for agreement, group in gt.set_index('imgName').loc[sample_id].groupby('agreement'):\n",
    "    color = colors[agreement - 1]\n",
    "    xs, ys = (group[['X', 'Y']].values * scale_f).astype(int).T\n",
    "    ax.plot(xs, ys, 'o', ms=25*scale_f, mec=color, mfc='none', mew=0.8)\n",
    "\n",
    "plt.savefig('figures/pnn-mr-sample.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_sr = pd.read_csv('data/perineuronal-nets/train/annotations.csv')\n",
    "\n",
    "sample_id = '034_B4_s06_C1.tif'\n",
    "img = imread('data/perineuronal-nets/train/fullFrames/' + sample_id)\n",
    "\n",
    "new_shape = (np.array(img.shape) / 10).astype(int)\n",
    "scale_f = new_shape[0] / img.shape[0]\n",
    "\n",
    "img = transform.resize(img, new_shape)\n",
    "\n",
    "p2, p98 = np.percentile(img, (0.1, 99.9))\n",
    "img = exposure.rescale_intensity(img, in_range=(p2, p98))\n",
    "\n",
    "img = matplotlib.cm.viridis(img)[:,:,:3]\n",
    "\n",
    "plt.figure(figsize=(14, 14))\n",
    "plt.imshow(img)\n",
    "ax = plt.gca()\n",
    "ax.set_axis_off()\n",
    "\n",
    "xy = gt_sr.set_index('imageName').loc[sample_id, ['X', 'Y']].values * scale_f\n",
    "xs, ys = xy.astype(int).T\n",
    "ax.plot(xs, ys, 'o', ms=35 * scale_f, mec='red', mfc='none', mew=0.8)\n",
    "\n",
    "plt.savefig('figures/pnn-sr-sample.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agreement between Raters in the Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "raters = gt.loc[0, 'AV':'VT'].index.values\n",
    "raters = np.array(raters).reshape(-1, 1)\n",
    "\n",
    "def agree(r1, r2):\n",
    "    a, b = gt[r1], gt[r2]\n",
    "    return jaccard_score(a, b)\n",
    "\n",
    "raters_agreement = pdist(raters, agree)\n",
    "raters_agreement = squareform(raters_agreement)\n",
    "\n",
    "mask = 1 - np.tri(len(raters), k=-1)\n",
    "\n",
    "raters_agreement = raters_agreement[1:, :-1]\n",
    "mask = mask[1:, :-1]\n",
    "\n",
    "ylabels = [f'R{i+2}' for i in range(len(raters)-1)]\n",
    "xlabels = [f'R{i+1}' for i in range(len(raters)-1)]\n",
    "sns.heatmap(raters_agreement, mask=mask, annot=True, square=True,\n",
    "            xticklabels=xlabels, yticklabels=ylabels, cmap='viridis',\n",
    "            cbar_kws=dict(location='left', label='Jaccard Index'))\n",
    "plt.savefig('figures/raters-agreement.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Total Cells counted by each Rater"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = gt.loc[:, 'AV':'VT'].sum(axis=0)\n",
    "counts.index = [f'R{i+1}' for i in range(len(counts))]\n",
    "counts = counts.to_frame(name='count')\n",
    "counts = counts.reset_index().rename({'index': 'rater'}, axis=1)\n",
    "ax = sns.barplot(data=counts, x='rater', y='count')\n",
    "\n",
    "mean = counts.mean().item()\n",
    "ax.axhline(mean, c='k', ls='--', lw=1.5)\n",
    "ax.set_ylim(1200, 1650)\n",
    "sns.despine()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stage 1: Localization/Counting Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_order = ('count/mare', 'count/game-3', 'pdet/f1_score')\n",
    "model_order = ('S-UNet', 'FRCNN', 'D-CSRNet')\n",
    "scorer_order = ('simple_regression', 'simple_classification', 'ordinal_regression', 'pairwise_balanced')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "runs = {\n",
    "    'S-UNet': list(Path('runs/experiment=perineuronal-nets/segmentation/').glob('unet_*')),\n",
    "    'FRCNN' : list(Path('runs/experiment=perineuronal-nets/detection/').glob('fasterrcnn_*')),\n",
    "    'D-CSRNet': list(Path('runs/experiment=perineuronal-nets/density/').glob('csrnet_*')),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def collect(model_name, run, csv_file):\n",
    "    run = Path(run)\n",
    "    cfg = OmegaConf.load(run / '.hydra' / 'config.yaml')\n",
    "    patch_size = cfg['data']['validation']['patch_size']\n",
    "\n",
    "    csv_path = run / 'test_predictions' / csv_file\n",
    "    if not csv_path.exists():\n",
    "        print(f'Skipping not found: {csv_path}')\n",
    "        return pd.DataFrame()\n",
    "    \n",
    "    data = pd.read_csv(csv_path, index_col=0)\n",
    "    data['model'] = model_name\n",
    "    data['patch_size'] = patch_size\n",
    "    \n",
    "    return data\n",
    "\n",
    "metrics = pd.concat([collect(k, r, 'all_metrics.csv.gz') for k, v in runs.items() for r in v], ignore_index=True)\n",
    "predictions = pd.concat([collect(k, r, 'all_gt_preds.csv.gz') for k, v in runs.items() for r in v], ignore_index=True)\n",
    "\n",
    "predictions['agreement'] = predictions['agreement'].fillna(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's the best patch size?\n",
    "Show trade-off between patch size and detection/counting performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "sns.set_theme(context='talk', style='ticks', font_scale=1.0)\n",
    "\n",
    "def compare_patch_sizes_plot(data, metric, metric_label, mode, fmt='.3f', ylim=(0,1), legend_bbta=(1,1)):\n",
    "    data = data.rename({'patch_size': 'Patch Size'}, axis=1)\n",
    "    g = sns.relplot(data=data, kind='line', col='model',\n",
    "                    x='thr', y=metric, hue='Patch Size', ci=None,\n",
    "                    facet_kws=dict(margin_titles=True, legend_out=True),\n",
    "                    aspect=1.2, height=4.5)\n",
    "\n",
    "    data = data.groupby(['model', 'Patch Size', 'thr']).mean()\n",
    "    best_points = data.groupby('model')[metric]\n",
    "    best_points = best_points.idxmin() if mode == 'min' else best_points.idxmax()\n",
    "\n",
    "    g.set(ylim=ylim, xlim=(0, 1))\n",
    "    g.set_titles(col_template=\"{col_name}\")\n",
    "    g.set_axis_labels(x_var='threshold', y_var=metric_label)\n",
    "    for model, ax in g.axes_dict.items():\n",
    "        ax.grid(True, which='major')\n",
    "        ax.grid(True, which='minor', ls='dotted')\n",
    "        ax.get_xaxis().set_minor_locator(ticker.AutoMinorLocator(2))\n",
    "        ax.get_yaxis().set_minor_locator(ticker.AutoMinorLocator(2))\n",
    "        ax.xaxis.set_major_formatter('{x:g}')\n",
    "\n",
    "        best_point = data.loc[best_points[model]]\n",
    "        _, patch_size, thr = best_point.name\n",
    "        value = best_point[metric]\n",
    "        print(f'[{metric}] {model} ps={patch_size} thr={thr} value={value:.2f}')\n",
    "\n",
    "        ax.plot([thr], [value], 'X', c='k', ms=9, mec='k', mfc='w')\n",
    "        xytext = (5, -3) if mode == 'min' else (5, 3)\n",
    "        va='top' if mode == 'min' else 'bottom'\n",
    "        ax.annotate(f'{value:{fmt}}', xy=(thr, value), xytext=xytext,\n",
    "                    textcoords='offset points', fontsize='small',\n",
    "                    va=va, ha='left')\n",
    "    \n",
    "    sns.move_legend(g, \"center right\", bbox_to_anchor=legend_bbta, frameon=False,\n",
    "                    labelspacing=0.25, fontsize='small', title_fontsize='small')\n",
    "    return g\n",
    "\n",
    "\n",
    "data = metrics[metrics.thr.between(0, 1)]\n",
    "\n",
    "compare_patch_sizes_plot(data, 'count/mae', 'MAE', 'min', fmt='.2f', ylim=(0, 200), legend_bbta=(.85, .5)) \\\n",
    "    .savefig('figures/pnn-mae.pdf', bbox_inches='tight')\n",
    "\n",
    "compare_patch_sizes_plot(data, 'count/mare', 'MARE', 'min', fmt='.1%', legend_bbta=(.85, .50)) \\\n",
    "    .savefig('figures/pnn-mare.pdf', bbox_inches='tight')\n",
    "\n",
    "compare_patch_sizes_plot(data, 'count/game-3', 'GAME(3)', 'min', fmt='.1f', ylim=(45, 200), legend_bbta=(.85, .5)) \\\n",
    "    .savefig('figures/pnn-game3.pdf', bbox_inches='tight')\n",
    "\n",
    "compare_patch_sizes_plot(data, 'pdet/f1_score', r'$F_1$-score', 'max', fmt='.1%', legend_bbta=(.85, .63))\\\n",
    "    .savefig('figures/pnn-f1-score.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# PR Curves\n",
    "sns.set_theme(context='talk', style='ticks')\n",
    "\n",
    "def plot_pr(data, label, color):\n",
    "    mean_pr = data.groupby('thr').mean().reset_index().sort_values('pdet/recall', ascending=False)\n",
    "    mean_recalls = mean_pr['pdet/recall'].values\n",
    "    mean_precisions = mean_pr['pdet/precision'].values\n",
    "    \n",
    "    aps = []\n",
    "    for group_key, img_group in data.groupby('imgName'):\n",
    "        img_group = img_group.reset_index().sort_values('pdet/recall', ascending=False)\n",
    "        recalls = img_group['pdet/recall'].values\n",
    "        precisions = img_group['pdet/precision'].values\n",
    "        average_precision = - np.sum(np.diff(recalls) * precisions[:-1])  # sklearn's ap\n",
    "        aps.append(average_precision)\n",
    "    \n",
    "    mean_ap = np.mean(aps)\n",
    "    plt.plot(mean_recalls, mean_precisions, label=f'{label} ({mean_ap:.1%})', color=color)\n",
    "\n",
    "\n",
    "data = metrics.copy()\n",
    "data.loc[data['pdet/recall'] == 0, 'pdet/precision'] = 1.0\n",
    "grid = sns.FacetGrid(data=data, hue='patch_size', col='model', height=4, xlim=(0,1), ylim=(0,1.05), aspect=1.2)\n",
    "grid.map_dataframe(plot_pr)\n",
    "grid.set_xlabels('Recall')\n",
    "grid.set_ylabels('Precision')\n",
    "grid.set_titles(col_template=\"{col_name}\")\n",
    "\n",
    "f_scores = np.linspace(0.1, 0.9, num=9)\n",
    "for ax in grid.axes.flatten():\n",
    "    ax.legend(title='Patch Size', loc='lower left', ncol=1, fontsize='x-small', title_fontsize='x-small')\n",
    "    \n",
    "    for i, f_score in enumerate(f_scores):\n",
    "        label_it = i % 2 != 0\n",
    "        ls = '-' if label_it else '--'\n",
    "        lw = 1 if label_it else 0.8\n",
    "        x = np.linspace(0.01, 1)\n",
    "        y = f_score * x / (2 * x - f_score)\n",
    "        l, = ax.plot(x[y >= 0], y[y >= 0], color='gray', alpha=0.2, ls=ls, lw=lw)\n",
    "        if label_it:\n",
    "            ax.annotate(r'$F_1$={0:0.1f}'.format(f_score), xy=(0.85, y[45] + 0.02), fontsize='xx-small')\n",
    "\n",
    "grid.savefig('figures/pnn-pr-curves.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Density-based Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "density_metrics = pd.concat([collect(k, r, 'dmap_metrics.csv.gz') for k, v in runs.items() for r in v], ignore_index=True)\n",
    "density_metrics.groupby(['model', 'patch_size']).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How's performance on different agreement levels?\n",
    "\n",
    "Show best counting metrics when practitioners choose different GT based on agreement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# common funcs\n",
    "from multiprocessing import Pool, cpu_count\n",
    "\n",
    "def _temp_func(args):\n",
    "    func, name, group = args\n",
    "    return func(group), name\n",
    "\n",
    "def applyParallel(dfGrouped, func):    \n",
    "    \n",
    "    with Pool(cpu_count()) as p:\n",
    "        gen = [(func, name, group) for name, group in dfGrouped]\n",
    "        ret_list = p.map(_temp_func, tqdm(gen))\n",
    "    \n",
    "    retLst, top_index = zip(*ret_list)\n",
    "    return pd.concat(retLst, keys=top_index)\n",
    "\n",
    "def _compute(x):\n",
    "    x = detection_and_counting(x, image_hw=(2000, 2000))\n",
    "    return pd.Series(x)\n",
    "\n",
    "\n",
    "def drop_empty_gp(x):\n",
    "    empty = x.X.isna() & x.Xp.isna()\n",
    "    return x[~empty]\n",
    "    \n",
    "\n",
    "def compute_metrics_by_agreement(data, grouping, parallel=True):\n",
    "    is_positive = ~data.Xp.isna() \n",
    "    \n",
    "    filtered = []\n",
    "    for i in range(1, 8):\n",
    "        tmp = data.copy()\n",
    "        tmp.loc[(tmp.agreement < i), 'X'] = None\n",
    "        tmp = drop_empty_gp(tmp)\n",
    "        tmp = tmp.assign(min_raters=i)\n",
    "        filtered.append(tmp)\n",
    "        \n",
    "    data = pd.concat(filtered, ignore_index=True)\n",
    "    data = data.groupby(grouping)\n",
    "    if parallel:\n",
    "        data = applyParallel(data, _compute)\n",
    "        data = data.unstack()\n",
    "        data.index.names = grouping\n",
    "        return data\n",
    "        \n",
    "    return data.progress_apply(_compute)\n",
    "\n",
    "def compute_ap(data):\n",
    "    pr = data.sort_values('pdet/recall', ascending=False)\n",
    "    recalls = pr['pdet/recall'].values\n",
    "    precisions = pr['pdet/precision'].values\n",
    "    ap = - np.sum(np.diff(recalls) * precisions[:-1])\n",
    "    return ap\n",
    "\n",
    "def build_map_table(data):\n",
    "    data = data.copy().reset_index()\n",
    "    model_grouper = ['model', 'patch_size']\n",
    "    if 'seed' in data.columns:\n",
    "        model_grouper.append('seed')\n",
    "        \n",
    "    if 'scorer' in data.columns:\n",
    "        model_grouper.append('scorer')\n",
    "        data['thr'] = data['re_thr_quantile']\n",
    "    \n",
    "    aps = data.groupby(model_grouper + ['min_raters', 'imgName']).apply(compute_ap)\n",
    "    mean_aps = aps.reset_index().groupby(model_grouper + ['min_raters']).mean()\n",
    "    mean_aps = mean_aps.rename(columns={0: 'mean_ap'})\n",
    "    return mean_aps\n",
    "\n",
    "def build_metrics_table(\n",
    "    data, \n",
    "    metric=['count/mare', 'count/game-3', 'pdet/f1_score'],\n",
    "    best_metric=None,\n",
    "    mode='min',\n",
    "    ci=False,\n",
    "    return_config=False\n",
    "):   \n",
    "\n",
    "    metric = metric if isinstance(metric, list) else [metric]\n",
    "    best_metric = metric if best_metric is None else best_metric\n",
    "    best_metric = best_metric if isinstance(best_metric, list) else [best_metric] * len(metric)\n",
    "    mode = mode if isinstance(mode, list) else [mode] * len(metric)\n",
    "    \n",
    "    assert len(metric) == len(best_metric), 'best_metric must be 1 or of the same size of metric'\n",
    "    assert len(metric) == len(mode), 'mode must be 1 or of the same size of metric'\n",
    "    \n",
    "    data = data.copy().reset_index()\n",
    "    model_grouper = ['model', 'patch_size']\n",
    "    if 'seed' in data.columns:\n",
    "        model_grouper.append('seed')\n",
    "        \n",
    "    if 'scorer' in data.columns:\n",
    "        model_grouper.append('scorer')\n",
    "        data['thr'] = data['re_thr_quantile']\n",
    "    \n",
    "    grouped = data.groupby(model_grouper + ['thr', 'min_raters'])\n",
    "    m, s = grouped.mean(), grouped.std()\n",
    "    \n",
    "    tables = []\n",
    "    configs = []\n",
    "    for metr, best_metr, mod in zip(metric, best_metric, mode):\n",
    "        best_points = m # if not ci else (m + s) if mod == 'min' else (m - s)\n",
    "        best_points = best_points.groupby(model_grouper + ['min_raters'])[best_metr]\n",
    "        best_points = best_points.idxmin() if mod == 'min' else best_points.idxmax()\n",
    "\n",
    "        table = m.loc[best_points, [metr]]\n",
    "        if ci:\n",
    "            table = table.combine(s.loc[best_points, [metr]], lambda x,y: x.combine(y, lambda w,z: (w,z)))\n",
    "        \n",
    "        table = table.reset_index().melt(id_vars=model_grouper + ['thr', 'min_raters'], var_name='metric')\n",
    "        tables.append(table)\n",
    "        configs.append(best_points)\n",
    "\n",
    "    table = pd.concat(tables)\n",
    "    \n",
    "    if return_config:\n",
    "        return table, configs\n",
    "    \n",
    "    return table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     4
    ]
   },
   "outputs": [],
   "source": [
    "# let's compute metrics\n",
    "p1_metrics = compute_metrics_by_agreement(predictions, ['model', 'patch_size', 'thr', 'imgName', 'min_raters'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cimax(args):\n",
    "    imax = args.map(lambda x: x[0] - x[1]*0).idxmax()\n",
    "    return args.loc[imax]\n",
    "\n",
    "def cimin(args):\n",
    "    imin = args.map(lambda x: x[0] + x[1]*0).idxmin()\n",
    "    return args.loc[imin]\n",
    "\n",
    "metr = ['count/mae', 'count/mare', 'count/game-3', 'pdet/f1_score']\n",
    "ci = True\n",
    "modes = ['min', 'min', 'min', 'max'] \n",
    "aggr = [cimin, cimin, cimin, cimax] if ci else modes\n",
    "\n",
    "p1_table, configs = build_metrics_table(p1_metrics, metric=metr, mode=modes, ci=ci, return_config=True)\n",
    "\n",
    "prec = 1\n",
    "xfm = {\n",
    "    'count/mae': lambda x: f'{x[0]:.1f};{x[1]:.1f}',\n",
    "    'count/mare': lambda x: f'{100*x[0]:.{prec}f};{100*x[1]:.{prec}f}',\n",
    "    'count/game-3': lambda x: f'{x[0]:.1f};{x[1]:.1f}',\n",
    "    'pdet/f1_score': lambda x: f'{100*x[0]:.{prec}f};{100*x[1]:.{prec}f}',\n",
    "} if ci else {\n",
    "    'count/mae': '{:.1f}'.format,\n",
    "    'count/mare': lambda x: f'{100*x:.1f}',\n",
    "    'count/game-3': '{:.1f}'.format,\n",
    "    'pdet/f1_score': lambda x: f'{100*x:.1f}', # '{:.0%}'.format,\n",
    "}\n",
    "\n",
    "aggr_per_metric = {k: v for k, v in zip(metr, aggr)}\n",
    "\n",
    "def take_best(a):\n",
    "    return a['value'].aggregate(aggr_per_metric[a.name[-1]])\n",
    "\n",
    "p1_table = p1_table.groupby(['model', 'min_raters', 'metric']).apply(take_best).rename('value')\n",
    "p1_table = p1_table.unstack('metric').transform(xfm).rename_axis('metric', axis=1).stack().rename('value')\n",
    "p1_table = p1_table.reset_index().pivot(index=['metric', 'model'], columns='min_raters', values='value')\n",
    "p1_table = p1_table.reindex(metr, level=0).reindex(model_order, level=1)\n",
    "\n",
    "print(p1_table.to_latex(escape=False, multirow=True))\n",
    "display(p1_table)\n",
    "\n",
    "p1_table = p1_table[[1,4,5,7]]\n",
    "\n",
    "print(p1_table.to_latex(escape=False, multirow=True))\n",
    "display(p1_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def find_best_worst_image(x, metric, mode):\n",
    "    x = x.reset_index()\n",
    "    m = x[metric]\n",
    "    return pd.Series({\n",
    "        'best' : x.loc[m.idxmin() if mode == 'min' else m.idxmax(), 'imgName'],\n",
    "        'worst': x.loc[m.idxmax() if mode == 'min' else m.idxmin(), 'imgName']\n",
    "    })\n",
    "\n",
    "best_metric = 'pdet/f1_score'\n",
    "best_metric_configs = configs[-1]\n",
    "mode = 'max'\n",
    "\n",
    "best_worst_images = p1_metrics \\\n",
    "    .reset_index().set_index(['model', 'patch_size', 'thr', 'min_raters']) \\\n",
    "    .loc[best_metric_configs] \\\n",
    "    .groupby(['model', 'patch_size', 'thr', 'min_raters']).apply(find_best_worst_image, best_metric, mode) \\\n",
    "    .mode()\n",
    "\n",
    "tmp = best_metric_configs.reset_index()\n",
    "selector = (((tmp.model == 'S-UNet')   & (tmp.patch_size == 320)) |\n",
    "            ((tmp.model == 'FRCNN')    & (tmp.patch_size == 640)) | \n",
    "            ((tmp.model == 'D-CSRNet') & (tmp.patch_size == 640)) )\n",
    "selector = selector & tmp.min_raters.isin([1, 7])\n",
    "best_configs = tmp[selector][best_metric].values\n",
    "best_configs\n",
    "\n",
    "indexed_preds = predictions.set_index(['model', 'patch_size', 'thr', 'imgName'])\n",
    "\n",
    "for img in ('best', 'worst'):\n",
    "    imgName = best_worst_images.loc[0, img]\n",
    "    image = imread('data/perineuronal-nets/test/fullFrames/' + imgName)\n",
    "    image = matplotlib.cm.viridis(image)[:,:,:3]\n",
    "    image = resize(image, (500, 500))\n",
    "    image = (255 * image).astype(np.uint8)\n",
    "    image = image[:250, 125:375, :]\n",
    "    \n",
    "    imsave(f'figures/{img}_clean.png', image)\n",
    "    \n",
    "    for model, patch_size, thr, min_raters in best_configs:\n",
    "        preds = indexed_preds.loc[(model, patch_size, thr, imgName)].reset_index().copy()\n",
    "        preds.loc[(preds.agreement < min_raters), ['X', 'Y']] = None\n",
    "        preds = drop_empty_gp(preds)\n",
    "        preds.loc[:, ['X', 'Y']] /= 4\n",
    "        preds.loc[:, ['Xp', 'Yp']] /= 4\n",
    "        \n",
    "        sel = ( (preds.X.isna() | (preds.X.between(125, 375) & preds.Y.between(0, 250))) | \n",
    "                (preds.Xp.isna() | (preds.Xp.between(125, 375) & preds.Yp.between(0, 250))) )\n",
    "        \n",
    "        preds = preds[sel]\n",
    "        preds = drop_empty_gp(preds)\n",
    "        preds.loc[:, 'X'] -= 125\n",
    "        preds.loc[:, 'Xp'] -= 125\n",
    "        \n",
    "        drawn = draw_groundtruth_and_predictions(image, preds, radius=5)\n",
    "        fname = f'figures/{img}_img_{model.lower()}_{patch_size}_raters_{min_raters}.png'\n",
    "        imsave(fname, drawn)\n",
    "        \n",
    "        gt_only_img = f'figures/{img}_gt_raters_{min_raters}.png'\n",
    "        if not Path(gt_only_img).exists():\n",
    "            gt_sel = ( (predictions.imgName == imgName)\n",
    "                     & (~predictions.agreement.isna())\n",
    "                     & (predictions.agreement >= min_raters)\n",
    "                     )\n",
    "            \n",
    "            gt_yx = predictions[gt_sel][['Y', 'X']].drop_duplicates().dropna()\n",
    "            gt_yx.loc[:, 'X'] = (gt_yx.X / 4) - 125\n",
    "            gt_yx.loc[:, 'Y'] = (gt_yx.Y / 4)\n",
    "            gt_yx = gt_yx[gt_yx.X.between(0, 250) & gt_yx.Y.between(0, 250)].values\n",
    "            \n",
    "            gt_only = draw_points(image, gt_yx, radius=5, marker='square', color=[255,255,0]) # YELLOW\n",
    "            imsave(gt_only_img, gt_only)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How much does rescoring (stage 2) increase performance?\n",
    "Compare counting and detection metrics of stage-1 only models and stage-2 refinement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runs_score_path = Path('runs_score')\n",
    "runs_score = {\n",
    "    #'AR': runs_score_path.glob('method=simple_regression,seed=*'),\n",
    "    #'AC': runs_score_path.glob('method=simple_classification,seed=*'),\n",
    "    #'OR': runs_score_path.glob('method=ordinal_regression,seed=*'),\n",
    "    #'RL': runs_score_path.glob('method=pairwise_balanced,seed=*'),\n",
    "    'Agreement Regression': runs_score_path.glob('method=simple_regression,seed=*'),\n",
    "    'Agreement Classification': runs_score_path.glob('method=simple_classification,seed=*'),\n",
    "    'Ordinal Regression': runs_score_path.glob('method=ordinal_regression,seed=*'),\n",
    "    'Rank Learning': runs_score_path.glob('method=pairwise_balanced,seed=*'),\n",
    "}\n",
    "\n",
    "def collect_scores(model_name, run):\n",
    "    run = Path(run)\n",
    "    csv_path = run / 'test_predictions' / 'all_gt_preds.csv.gz'\n",
    "    data = pd.read_csv(csv_path, index_col=0)\n",
    "    data['model'] = model_name\n",
    "    data['seed'] = int(run.name.split('=')[-1])\n",
    "    return data\n",
    "\n",
    "score_data = [collect_scores(k, run) for k, runs in runs_score.items() for run in runs]\n",
    "score_data = pd.concat(score_data, ignore_index=True)\n",
    "\n",
    "test_images = score_data.groupby('seed').imgName.unique().to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# best configs for maximum recall\n",
    "\n",
    "def max_recall(data):\n",
    "    data = data.sort_values(['pdet/recall', 'pdet/precision'], ascending=[False, False])\n",
    "    return data.head(1).index.values\n",
    "\n",
    "p1_metrics.xs(1, level='min_raters').groupby(['model', 'patch_size', 'thr']).mean().groupby('model').apply(max_recall)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Samples per Scorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = PerineuronalNetsRankDataset(mode='patches')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_theme(context='notebook', style='ticks', font_scale=1)\n",
    "\n",
    "p2i = dataset.annot.reset_index().set_index(['imgName','X','Y'])\n",
    "\n",
    "so = ('Pair-wise Regression', 'Ordinal Regression', 'Agreement Classification', 'Agreement Regression')\n",
    "\n",
    "sample_idx = rank_data.groupby(['model', 'agreement'])\\\n",
    "    .apply(lambda x: x.nlargest(10, 'score')).droplevel(-1)\\\n",
    "    .apply(lambda x: p2i.loc[tuple(x[['imgName', 'X', 'Y']].values), 'index'], axis=1)\\\n",
    "\n",
    "nr=1\n",
    "fig, axes = plt.subplots(7, len(so), figsize=(17,4))\n",
    "for i, scorer in enumerate(so):\n",
    "    axes[0, i].set_title(scorer)\n",
    "    for j, agreement in enumerate(range(7, 0, -1)):\n",
    "        samples = sample_idx.loc[(scorer, agreement)]\n",
    "        cell_images = [dataset[i][0] for i in samples]\n",
    "        cell_images = [matplotlib.cm.viridis(c) for c in cell_images]\n",
    "        cell_images = np.stack(cell_images)[:,:,:,:3]\n",
    "        image = montage(cell_images, grid_shape=(nr, len(cell_images) / nr), padding_width=5, fill=(1, 1, 1), multichannel=True)\n",
    "        axes[j, i].imshow(image)\n",
    "        axes[j, i].set_axis_off()\n",
    "\n",
    "for j, agreement in enumerate(range(7, 0, -1)):\n",
    "    axes[j, 0].set_ylabel(str(agreement))\n",
    "    \n",
    "plt.subplots_adjust(wspace=0.1, hspace=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stage-1 Only Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get best patch_size per method, all thresholds\n",
    "tmp = p1_metrics.reset_index()\n",
    "\n",
    "selector = (((tmp.model == 'S-UNet')   & (tmp.patch_size == 320)) |\n",
    "            ((tmp.model == 'FRCNN')    & (tmp.patch_size == 640)) | \n",
    "            ((tmp.model == 'D-CSRNet') & (tmp.patch_size == 640)) )\n",
    "\n",
    "tmp = tmp[selector]\n",
    "\n",
    "# keep only test set\n",
    "tmp = pd.concat([tmp[tmp.imgName.isin(images)].assign(seed=seed) for seed, images in test_images.items()], ignore_index=True)\n",
    "p1_test_metrics = tmp.set_index(p1_metrics.index.names + ['seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1t_table = build_metrics_table(p1_test_metrics, metric=metr, mode=modes)\n",
    "p1t_table = p1t_table.groupby(['model', 'patch_size', 'min_raters', 'metric']).value.aggregate(['mean', 'std'])\n",
    "p1t_table = p1t_table.unstack('metric')\n",
    "\n",
    "pct_f = lambda x: f'{100*x:.1f}'\n",
    "flo_f = '{:.1f}'.format\n",
    "p1t_table = p1t_table.transform({\n",
    "    ('mean', 'count/mae' ): flo_f,\n",
    "    ('std' , 'count/mae' ): flo_f,\n",
    "    ('mean', 'count/mare'   ): pct_f,\n",
    "    ('std' , 'count/mare'   ): pct_f,\n",
    "    ('mean', 'count/game-3' ): flo_f,\n",
    "    ('std' , 'count/game-3' ): flo_f,\n",
    "    ('mean', 'pdet/f1_score'): pct_f,\n",
    "    ('std' , 'pdet/f1_score'): pct_f,\n",
    "})\n",
    "\n",
    "p1t_table = p1t_table['mean'] + ' $\\pm$ ' + p1t_table['std']\n",
    "p1t_table = p1t_table.rename_axis('metric', axis=1).stack().rename('value').reset_index()\n",
    "p1t_table = p1t_table.pivot(index=['metric', 'model'], columns='min_raters', values='value')\n",
    "p1t_table = p1t_table.reindex(metr, level=0).reindex(model_order, level=1)\n",
    "p1t_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(p1_table, p1t_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1_map_table = build_map_table(p1_metrics)\\\n",
    "    .unstack('min_raters')\\\n",
    "    .applymap(lambda x: f'{100*x:.1f}')\n",
    "\n",
    "p1t_map_table = build_map_table(p1_test_metrics)\\\n",
    "    .reset_index()\\\n",
    "    .groupby(['model', 'patch_size', 'min_raters'])\\\n",
    "    .mean_ap.aggregate(['mean', 'std'])\\\n",
    "    .applymap(lambda x: f'{100*x:.1f}')\n",
    "\n",
    "p1t_map_table = p1t_map_table['mean'] + ' $\\pm$ ' + p1t_map_table['std']\n",
    "p1t_map_table = p1t_map_table\\\n",
    "    .rename('value')\\\n",
    "    .reset_index()\\\n",
    "    .pivot(index=['model', 'patch_size'], columns='min_raters', values='value')\n",
    "\n",
    "display(p1_map_table, p1t_map_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Score vs Agreement Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get best config per model, maximum recall\n",
    "selector = (((predictions.model == 'S-UNet')   & (predictions.patch_size == 320) & (predictions.thr == 0.1)) |\n",
    "            ((predictions.model == 'FRCNN')    & (predictions.patch_size == 640) & (predictions.thr == 0.0)) | \n",
    "            ((predictions.model == 'D-CSRNet') & (predictions.patch_size == 640) & (predictions.thr == 0.0)))\n",
    "\n",
    "# keep only test sets\n",
    "keep = np.unique(np.concatenate(list(test_images.values()))).tolist()\n",
    "selector = selector & predictions.imgName.isin(keep)\n",
    "\n",
    "p1_data = predictions[selector].copy()\n",
    "p1_data['agreement'] = p1_data['agreement'].fillna(0)\n",
    "p1_data['seed'] = 23\n",
    "\n",
    "p2_data = score_data.copy()\n",
    "p2_data['patch_size'] = -1\n",
    "\n",
    "rdata = pd.concat([p1_data, p2_data], ignore_index=True)\n",
    "\n",
    "def normalize_scores(data):\n",
    "    data['score'] = StandardScaler().fit_transform(data['score'].values.reshape(-1, 1)) # * 0.5 + 0.5\n",
    "    return data\n",
    "\n",
    "rdata = rdata.groupby(['model', 'patch_size', 'seed']).apply(normalize_scores)\n",
    "rdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     10,
     53
    ],
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sns.set_theme(context='talk', style='ticks', font_scale=1.0)\n",
    "\n",
    "plot_data = rdata[~rdata.score.isna() & (rdata.agreement > 0)].copy()\n",
    "plot_data['agreement'] = plot_data.agreement.astype(int)\n",
    "plot_data = plot_data[['score', 'agreement', 'model']]\n",
    "\n",
    "order = [\n",
    "    'S-UNet',\n",
    "    'FRCNN',\n",
    "    'D-CSRNet',\n",
    "    'Agreement Regression',\n",
    "    'Agreement Classification',\n",
    "    'Ordinal Regression',\n",
    "    'Rank Learning',\n",
    "    #'AR',\n",
    "    #'AC',\n",
    "    #'OR',\n",
    "    #'RL',\n",
    "]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16, 6))\n",
    "width = 0.8\n",
    "sns.boxenplot(data=plot_data, y='score', x='model', hue='agreement', order=order, palette='rocket', ax=ax, width=width, showfliers=False)\n",
    "ax.set_ylim([-3.5, 2.5])\n",
    "ax.set_yticks(range(-3, 3))\n",
    "ax.set_yticklabels(range(-3, 3))    \n",
    "ax.axhline(xmax=.95, c='k', zorder=-10, lw=1.5)\n",
    "\n",
    "def corr_coeff(data, **kws):\n",
    "    sel = (~data.score.isna()) & (~data.agreement.isna())\n",
    "    x = data.loc[sel, 'score']\n",
    "    y = data.loc[sel, 'agreement']\n",
    "    r, p = scipy.stats.pearsonr(x, y)\n",
    "    return r\n",
    "\n",
    "def lin_fit(data, **kws):\n",
    "    sel = (~data.score.isna()) & (~data.agreement.isna())\n",
    "    \n",
    "    p = []\n",
    "    grouped = data[sel].groupby('agreement')\n",
    "    min_num = grouped.model.count().min()\n",
    "    for _ in range(50):\n",
    "        y, x = grouped.sample(min_num)[['score', 'agreement']].values.T\n",
    "        z = np.polyfit(x, y, 1)   \n",
    "        p.append(z)\n",
    "    \n",
    "    p = np.mean(p, axis=0)\n",
    "    p = np.poly1d(p)\n",
    "    return p\n",
    "\n",
    "grouped = plot_data.groupby('model')\n",
    "corrs = grouped.apply(corr_coeff)\n",
    "linfits = grouped.apply(lin_fit)\n",
    "\n",
    "display(corrs)\n",
    "\n",
    "labels = [l.get_text() for l in ax.get_xticklabels()]\n",
    "labels = ['{}\\n$r$={:.2f}'.format(l.replace(' ', '\\n'), corrs[l]) for l in labels]\n",
    "\n",
    "ax.set_xticklabels(labels)\n",
    "ax.set_xlabel(None)\n",
    "ax.set_ylabel('z-score')\n",
    "\n",
    "ax.grid(which='major', axis='y', ls='-', lw=.75, zorder=-10)\n",
    "ax.tick_params(axis='x', color='white', labelbottom=False, labeltop=True)\n",
    "\n",
    "offsets = (0.18, 0.18, 0.1, 0.1, 0.1, 0.1, 0.1)\n",
    "for label, offset in zip(ax.get_xticklabels(), offsets):\n",
    "    label.set_y(label.get_position()[1] - offset)\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "\n",
    "legend_order = (1, 2, 3, 4, 5, 6, 7)\n",
    "handles = [handles[i-1] for i in legend_order]\n",
    "labels = [labels[i-1] for i in legend_order]\n",
    "\n",
    "ax.legend(handles, labels, title='agreement',\n",
    "          ncol=7, loc='lower right', bbox_to_anchor=(1, 0), # bbox_to_anchor=(.12,.12,.77,1),\n",
    "          fontsize='x-small', title_fontsize='x-small',\n",
    "          labelspacing=0.2, columnspacing=1, framealpha=1, fancybox=False)\n",
    "\n",
    "for i, line in enumerate(linfits[order]):\n",
    "    x = [i - 4 * width / 7, i + 4 * width / 7]\n",
    "    y = line([0, 8])\n",
    "    ax.plot(x, y, c='w', ls='--', path_effects=[pe.Stroke(linewidth=4, foreground='k'), pe.Normal()])\n",
    "\n",
    "sns.despine(bottom=True)\n",
    "plt.savefig('figures/score-vs-agreement.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_theme(context='notebook', style='ticks', font_scale=1.5)\n",
    "\n",
    "plot_data = rdata.fillna({'agreement': 0, 'score': -100000}).groupby(['model', 'X', 'Y', 'agreement']).score.mean().reset_index()\n",
    "sorted_samples = plot_data.groupby('model')\\\n",
    "    .apply(lambda x: \\\n",
    "           x.sort_values(['score', 'Y', 'X'], ascending=[False, True, True]).agreement)\n",
    "sorted_samples = sorted_samples.droplevel(-1).reset_index()\n",
    "sorted_samples['model'] = sorted_samples['model'].str.replace(' ', '\\n')\n",
    "\n",
    "order = [\n",
    "    'S-UNet',\n",
    "    'FRCNN',\n",
    "    'D-CSRNet',\n",
    "    'Agreement\\nRegression',\n",
    "    'Agreement\\nClassification',\n",
    "    'Ordinal\\nRegression',\n",
    "    'Rank\\nLearning',\n",
    "]\n",
    "\n",
    "def heatmap_plot(data, color, **kws):\n",
    "    data = data.agreement.values\n",
    "    rows = 32\n",
    "    pad = (-data.size) % rows\n",
    "    \n",
    "    mask = np.zeros_like(data)\n",
    "    mask = np.pad(mask, (0, pad), constant_values=1).reshape(-1, rows)\n",
    "    data = np.pad(data, (0, pad), constant_values=100).reshape(-1, rows)\n",
    "    \n",
    "    sns.heatmap(ax=plt.gca(), data=data, mask=mask, **kws)\n",
    "    \n",
    "    \n",
    "g = sns.FacetGrid(data=sorted_samples, col='model', aspect=.45, height=5, col_order=order)\n",
    "cbar_ax = g.fig.add_axes([.99, .2, .01, .60])  # create a colorbar axes\n",
    "\n",
    "g = g.map_dataframe(heatmap_plot, vmin=0, vmax=7, square=True, antialiased=True, rasterized=True,\n",
    "                    cbar_ax=cbar_ax, cbar_kws=dict(\n",
    "                        ticks=range(8),\n",
    "                        ticklocation='right', orientation='vertical',\n",
    "                        label='agreement',\n",
    "                    ))\n",
    "for ax in g.axes.flatten():\n",
    "    ax.axis('off')\n",
    "    \n",
    "g.set_titles(col_template=\"{col_name}\")\n",
    "g.tight_layout()\n",
    "g.fig.subplots_adjust(wspace=.05)#, hspace=0.05)\n",
    "g.savefig('figures/score-gradient.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stage-2 Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rescored_runs = {\n",
    "    'S-UNet': [\n",
    "        ('runs/experiment=perineuronal-nets/segmentation/unet_320', 0.1),\n",
    "    ],\n",
    "    'FRCNN' : [\n",
    "        ('runs/experiment=perineuronal-nets/detection/fasterrcnn_640', 0.00),\n",
    "    ],\n",
    "    'D-CSRNet': [\n",
    "        ('runs/experiment=perineuronal-nets/density/csrnet_640', 0.00),\n",
    "    ],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def collect_rescored(model_name, run, thr):\n",
    "    run = Path(run)\n",
    "    cfg = OmegaConf.load(run / '.hydra' / 'config.yaml')\n",
    "    patch_size = cfg['data']['validation']['patch_size']\n",
    "\n",
    "    preds = []\n",
    "    csv_paths = (run / 'test_predictions').glob('all_gt_preds_rescored_*seed*_imgsplit.csv.gz')\n",
    "    for csv_path in csv_paths:\n",
    "        method_and_seed = csv_path.name[len('all_gt_preds_rescored_'):-len('_imgsplit.csv.gz')]\n",
    "        rescore_method, seed = method_and_seed.split('-')\n",
    "        seed = int(seed[len('seed'):])\n",
    "        \n",
    "        data = pd.read_csv(csv_path)\n",
    "        data = data[(data.thr == thr) & (data.imgName.isin(test_images[seed]))]\n",
    "        data['model'] = model_name\n",
    "        data['patch_size'] = patch_size\n",
    "        data['scorer'] = rescore_method\n",
    "        data['seed'] = seed\n",
    "        \n",
    "        preds.append(data)\n",
    "    \n",
    "    return pd.concat(preds, ignore_index=True)\n",
    "\n",
    "rescored_predictions = pd.concat([collect_rescored(k, r, t) for k, v in rescored_runs.items() for r, t in v], ignore_index=True)\n",
    "rescored_predictions['agreement'] = rescored_predictions.agreement.fillna(0)\n",
    "\n",
    "\n",
    "def apply_percentile_thresholds(gp):\n",
    "    quantiles = np.linspace(0, 1, 201)\n",
    "    if gp.scorer.iloc[0] == 'no_rescore':\n",
    "        re_thrs = quantiles.tolist()\n",
    "    else:\n",
    "        re_thrs = gp.rescore.quantile(quantiles).tolist()\n",
    "    \n",
    "    quantiles = quantiles.tolist()\n",
    "    quantiles.append(2.)\n",
    "    re_thrs.append(re_thrs[-1] + 1)\n",
    "    \n",
    "    all_thresholded = []\n",
    "    for re_thr, q in zip(re_thrs, quantiles):\n",
    "        thresholded = gp.copy()\n",
    "        thresholded.loc[(gp.rescore < re_thr) | gp.rescore.isna(), 'Xp'] = None\n",
    "        thresholded = thresholded[~(thresholded.X.isna() & thresholded.Xp.isna())]\n",
    "        thresholded['re_thr'] = re_thr\n",
    "        thresholded['re_thr_quantile'] = q\n",
    "        all_thresholded.append(thresholded)\n",
    "    \n",
    "    return pd.concat(all_thresholded, ignore_index=True)    \n",
    "\n",
    "rescored_predictions = rescored_predictions.groupby(['patch_size', 'model', 'seed', 'thr', 'scorer'])\\\n",
    "                                           .progress_apply(apply_percentile_thresholds)\\\n",
    "                                           .reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p2_metrics = compute_metrics_by_agreement(\n",
    "    rescored_predictions,\n",
    "    ['model', 'patch_size', 'scorer', 're_thr_quantile', 'imgName', 'min_raters', 'seed']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1t_mean_ap = build_map_table(p1_test_metrics)\n",
    "p2_mean_ap = build_map_table(p2_metrics)\n",
    "\n",
    "#display(p1t_mean_ap, p2_mean_ap)\n",
    "tmp_p1 = pd.concat({'-': p1t_mean_ap}, names=['scorer']).reset_index().set_index(p2_mean_ap.index.names)\n",
    "combined = pd.concat((p2_mean_ap, tmp_p1))\\\n",
    "    .groupby(['model', 'patch_size', 'scorer', 'min_raters']).mean()\n",
    "\n",
    "diff = combined - combined.xs('-', level=2)\n",
    "\n",
    "def fmt(absolute, difference):\n",
    "    return f'{absolute:.2f} ({difference:.2f})'\n",
    "\n",
    "def styling(x):\n",
    "    diff = float(x.split(' ')[1].strip('()'))\n",
    "    color = '#ADFFAD' if diff > 0 else '#ffadad'if diff < 0 else 'none'\n",
    "    return f'background-color: {color}'\n",
    "\n",
    "combined.combine(diff, lambda x, y: x.combine(y, fmt)) \\\n",
    "    .reindex(model_order, level=0).reindex(('-',) + scorer_order, level=2) \\\n",
    "    .unstack('min_raters') \\\n",
    "    .style.applymap(styling)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "metr = ['count/mae', 'count/mare', 'count/game-3', 'pdet/f1_score']\n",
    "modes = ['min', 'min', 'min', 'max'] \n",
    "# modes = ['min', 'min', 'max'] \n",
    "\n",
    "p1t_table = build_metrics_table(p1_test_metrics, metric=metr, mode=modes).assign(scorer='-')\n",
    "p2_table = build_metrics_table(p2_metrics, metric=metr, mode=modes)\n",
    "\n",
    "#p1t_table = pd.concat({'-': p1t_table}, names=['scorer']).reset_index().set_index(p2_table.index.names)\n",
    "\n",
    "combined = pd.concat((p1t_table, p2_table), ignore_index=True)\\\n",
    "    .groupby(['model', 'patch_size', 'scorer', 'min_raters', 'metric']).value.mean().rename('value')\\\n",
    "    .reset_index().set_index(['metric', 'model', 'patch_size', 'scorer', 'min_raters'])\n",
    "\n",
    "#combined\n",
    "diff = combined - combined.xs('-', level='scorer')\n",
    "# display(combined, combined.xs('-', level='scorer'))\n",
    "\n",
    "def fmt(absolute, difference):\n",
    "    return f'{absolute:.2f} ({difference:.2f})'\n",
    "\n",
    "def styling_up(x):\n",
    "    diff = float(x.split(' ')[1].strip('()'))\n",
    "    color = '#ADFFAD' if diff > 0 else '#ffadad'if diff < 0 else 'none'\n",
    "    return f'background-color: {color}'\n",
    "\n",
    "def styling_down(x):\n",
    "    diff = float(x.split(' ')[1].strip('()'))\n",
    "    color = '#ADFFAD' if diff < 0 else '#ffadad'if diff > 0 else 'none'\n",
    "    return f'background-color: {color}'\n",
    "\n",
    "styles = {\n",
    "    'count/mae': styling_down,\n",
    "    'count/mare': styling_down,\n",
    "    'count/game-3': styling_down,\n",
    "    'pdet/f1_score': styling_up\n",
    "}\n",
    "\n",
    "def styling(x):\n",
    "    style_func = styles[x.name[0]]\n",
    "    return [style_func(i) for i in x.values]\n",
    "\n",
    "table = combined.combine(diff, lambda x, y: x.combine(y, fmt)) \\\n",
    "    .reindex(model_order, level='model')\\\n",
    "    .reindex(('-',) + scorer_order, level='scorer') \\\n",
    "    .unstack('min_raters')\\\n",
    "    .style.apply(styling, axis=1)\n",
    "\n",
    "display(table)\n",
    "\n",
    "\n",
    "def latex_fmt(a, d):\n",
    "    return f'{a:.2f}\\diff{{{d:.2f}}}'\n",
    "\n",
    "table = combined.combine(diff, lambda x, y: x.combine(y, latex_fmt))\\\n",
    "    .reindex(model_order, level='model')\\\n",
    "    .reindex(('-',) + scorer_order, level='scorer') \\\n",
    "    .unstack('min_raters')\\\n",
    "    .droplevel('patch_size', axis=0)\\\n",
    "    .droplevel(0, axis=1)\\\n",
    "    .loc['count/mae', [1, 4, 5,7]]\\\n",
    "    .rename({\n",
    "        '-': '',\n",
    "        'simple_regression': 'AR', #'Agreement Regression',\n",
    "        'simple_classification': 'AC', #'Agreement Classification',\n",
    "        'ordinal_regression': 'OR', #'Ordinal Regression',\n",
    "        'pairwise_balanced': 'RL', #'Rank Learning',\n",
    "    }, axis=0, level='scorer')\n",
    "    \n",
    "\n",
    "table = table.set_index(table.index.map(lambda x: x[0] + (' + ' + x[1] if x[1] else '')))\n",
    "print(table.to_latex(escape=False))\n",
    "table"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}