scripts/Figure_2A_4AB_CP_OM_Compare.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Cellpose vs Omnipose performance\n",
    "This script runs Cellpose and Omnipose on various 2D datasets: bact_phase, bact_fluor, cyto2, and worm.\n",
    "It is very long because I would like to reduce the repetition the loading and analysis code (factoring it into functions is overkill and kind of infeasible, since there are some special cases sprinkled throughout). I suggest using the ToC to navigate. \n",
    "\n",
    "Extended Data Figures 4,5 look at the failure modes of Cellpose. Figure 3 compares Omnipose to Cellpose on phase. Fig. 4 compares Omnipose to Cellpose on fluorescence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "from cellpose_omni import plot, models, io, dynamics\n",
    "from omnipose import utils\n",
    "import skimage.io\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('dark_background')\n",
    "import matplotlib as mpl\n",
    "%matplotlib inline\n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "import time, os, sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define datasets and models \n",
    "We will run the same code on each of the following datasets. We begin by defining names (used for saving output) as well as the paths to the test datasets. Later we will define the paths to the models trained ont he corresponding training datasets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = ['bact_phase','bact_fluor','cyto2','worm']\n",
    "\n",
    "basedirs = ['/home/kcutler/DataDrive/omnipose_all/phase/test_sorted',\n",
    "            '/home/kcutler/DataDrive/omnipose_all/fluor/test_sorted',\n",
    "            '/home/kcutler/DataDrive/cyto2/test',\n",
    "            '/home/kcutler/DataDrive/omnipose_train/worm_combined/test',\n",
    "           ]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "which = 0 # select which of the above models to run\n",
    "clean = 0 # control whether or not to re-run; still want to define the name of each experiment \n",
    "dataset = datasets[which]\n",
    "basedir = basedirs[which]\n",
    "\n",
    "mask_filter = '_masks'\n",
    "\n",
    "if dataset=='cyto2':\n",
    "    img_filter = '_img'\n",
    "else:\n",
    "    img_filter = ''\n",
    "    \n",
    "img_names = io.get_image_files(basedir,mask_filter,img_filter,look_one_level_down=True)\n",
    "mask_names = io.get_label_files(img_names, mask_filter,img_filter)\n",
    "# Obviously sorting the masks by corresponding base name is critical. The different suffixes mess this up, but sorted()\n",
    "# can take a function like the one below as a key. Modifying the default suffix requires the lambda syntax (e.g. cellpose). \n",
    "def getname(path,suffix='_masks'):\n",
    "    return os.path.splitext(Path(path).name)[0].replace(suffix,'')\n",
    "\n",
    "mask_names = sorted(mask_names,key=getname)\n",
    "img_names = sorted(img_names,key=getname)\n",
    "imgs = [skimage.io.imread(f) for f in img_names]\n",
    "masks_gt = [skimage.io.imread(f) for f in mask_names]\n",
    "\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count masks\n",
    "cnt = 0;\n",
    "for mask in masks_gt:\n",
    "    lbls = np.unique(utils.format_labels(mask))\n",
    "    cnt += len(lbls[lbls>0])\n",
    "print(cnt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize models\n",
    "The list of models I want to run depends on the dataset and the figure. The `bact_phase` dataset is the main dataset for the paper. The `bact_fluor` dataset is evaluated with Omnipose only because we are only interested in how Omnipose runs on it compared to the `bact_phase` dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cellpose_omni import core, models\n",
    "net_avg=False\n",
    "\n",
    "use_GPU = core.use_gpu()\n",
    "print('>>> GPU activated? %d'%use_GPU)\n",
    "\n",
    "\n",
    "if dataset == 'bact_phase':\n",
    "\n",
    "    model = [models.CellposeModel(gpu=use_GPU, model_type='bact_phase_cp'),\n",
    "             models.CellposeModel(gpu=use_GPU, model_type='bact_phase_omni')]\n",
    "    chans = [0,0]\n",
    "    \n",
    "elif dataset == 'bact_fluor':\n",
    "\n",
    "    model = [models.CellposeModel(gpu=use_GPU, model_type='bact_fluor_cp'),\n",
    "             models.CellposeModel(gpu=use_GPU, model_type='bact_fluor_omni')]\n",
    "    chans = [0,0]\n",
    "    \n",
    "elif dataset=='cyto2':\n",
    "\n",
    "    model = [models.CellposeModel(gpu=use_GPU, model_type='cyto2'),\n",
    "             models.CellposeModel(gpu=use_GPU, model_type='cyto2_omni')] # needs to be updated to use my bit depth version \n",
    "    chans = [[1,2],[2,1]] # this should be updated too with the new model and different training parameters \n",
    "    \n",
    "elif dataset=='worm':   \n",
    "    # pure worm cellpose models and omnipose as well as a combined bact+worm omnipose model\n",
    "    # all trained for 3201 epochs \n",
    "    model = [models.CellposeModel(gpu=use_GPU, model_type='worm_cp'),\n",
    "             models.CellposeModel(gpu=use_GPU, model_type='worm_omni'),\n",
    "             models.CellposeModel(gpu=use_GPU, model_type='worm_bact_omni')]\n",
    "\n",
    "    chans = [0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from cellpose_omni.io import logger_setup\n",
    "logger,log_file=logger_setup()\n",
    "resample = True\n",
    "transparency = True # save flows with transparency\n",
    "verbose = False\n",
    "cluster = True #use clustering algorithm for Omnipose \n",
    "\n",
    "imglist = imgs[:]\n",
    "N = len(imgs)\n",
    "names_phase = ['Cellpose','Mixed','Omnipose']\n",
    "print('datasest is {} consisting of {} images and {} labels. Clean flag is {}'.format(dataset,N,cnt,clean))\n",
    "if dataset == 'bact_phase':\n",
    "    names = names_phase\n",
    "    omni = [False,True,True]\n",
    "    n = len(names)\n",
    "    masks,flows,styles,d = [[]]*n,[[]]*n,[[]]*n,[[]]*n\n",
    "\n",
    "    if clean:\n",
    "\n",
    "        # Cellpose 'bact', no omni reconstruction \n",
    "        masks[0], flows[0], styles[0] = model[0].eval(imglist,channels=chans,rescale=None,mask_threshold=-1,flow_threshold=0,\n",
    "                                                      omni=omni[0],resample=resample,tile=False, transparency=transparency, \n",
    "                                                      verbose=verbose)\n",
    "        # Cellpose + omni ('Hybrid')\n",
    "        masks[1], flows[1], styles[1] = model[0].eval(imglist,channels=chans,rescale=None,mask_threshold=-1,flow_threshold=0,\n",
    "                                                      omni=omni[1],cluster=cluster,resample=resample,tile=False,transparency=transparency, \n",
    "                                                      verbose=verbose)\n",
    "        #Omnipose\n",
    "        masks[2], flows[2], styles[2] = model[1].eval(imglist,channels=chans,rescale=None,mask_threshold=-1,flow_threshold=0,\n",
    "                                                      omni=omni[2],cluster=cluster,resample=resample,tile=False,transparency=transparency, \n",
    "                                                      verbose=verbose)\n",
    "\n",
    "if dataset == 'bact_fluor':\n",
    "    names = ['Cellpose','Omnipose']\n",
    "    n = len(names)\n",
    "    masks,flows,styles,d = [[]]*n,[[]]*n,[[]]*n,[[]]*n\n",
    "    if clean: \n",
    "        cluster = False # both models perform better without clustering in this case\n",
    "        \n",
    "        for k in range(len(model)):\n",
    "            masks[k], flows[k], _ = model[k].eval(imglist,channels=chans,rescale=None,mask_threshold=-1,flow_threshold=-1,\n",
    "                                                  omni=k, cluster=cluster, resample=resample, tile=False, transparency=transparency, \n",
    "                                                  verbose=verbose)\n",
    "\n",
    "\n",
    "if dataset=='cyto2':\n",
    "    names = ['Cellpose','Mixed','Omnipose']\n",
    "    omni = [False,True,True]\n",
    "    n = len(names)\n",
    "    masks,flows,styles,d = [[]]*n,[[]]*n,[[]]*n,[[]]*n\n",
    "    if clean:\n",
    "        resample = False # both models perform better without resampling in this case\n",
    "        diam_threshold = 30 # I specified this at one point, clustering is on anyway\n",
    "        \n",
    "        diameters = [omnipose.core.diameters(mask) for mask in masks_gt]\n",
    "        \n",
    "        # cyto2\n",
    "        masks[0], flows[0], styles[0] = model[0].eval(imglist, channels=chans[0], diameter=diameters, mask_threshold=0, flow_threshold=0,\n",
    "                                                      omni=omni[0], resample=resample, tile=False, transparency=transparency, \n",
    "                                                      verbose=verbose)\n",
    "        # cyto2 + omni=True ('Mixed')\n",
    "        masks[1], flows[1], styles[1] = model[0].eval(imglist, channels=chans[0], diameter=diameters, mask_threshold=0, flow_threshold=0,\n",
    "                                                      omni=omni[1], cluster=cluster, resample=resample, tile=False, transparency=transparency, \n",
    "                                                      verbose=verbose, diam_threshold=diam_threshold)\n",
    "        # cyto2_omni\n",
    "        masks[2], flows[2], styles[2] = model[1].eval(imglist, channels=chans[1], diameter=diameters, mask_threshold=-1, flow_threshold=0,\n",
    "                                                      omni=omni[2], cluster=cluster, resample=resample, tile=False, transparency=transparency, \n",
    "                                                      verbose=verbose, diam_threshold=diam_threshold)\n",
    "\n",
    "\n",
    "if dataset=='worm':\n",
    "    names = ['Cellpose','Mixed','Omnipose','Omnipose_worm+bact']\n",
    "    omni = [False,True,True,True]\n",
    "    n = len(names)\n",
    "    masks,flows,styles,d = [[]]*n,[[]]*n,[[]]*n,[[]]*n\n",
    "    if clean:\n",
    "        # worm_cp, no omni reconstruction \n",
    "        masks[0], flows[0], styles[0] = model[0].eval(imglist, channels=chans, rescale=None, mask_threshold=-1, flow_threshold=0,\n",
    "                                                      omni=omni[0],cluster=cluster, resample=resample, tile=False, transparency=transparency, verbose=verbose)\n",
    "        # worm_cp + omni=True ('Mixed')\n",
    "        masks[1], flows[1], styles[1] = model[0].eval(imglist, channels=chans, rescale=None, mask_threshold=-1, flow_threshold=0,\n",
    "                                                      omni=omni[1],cluster=cluster, resample=resample, tile=False, transparency=transparency, verbose=verbose)\n",
    "        # worm_omni\n",
    "        masks[2], flows[2], styles[2] = model[1].eval(imglist, channels=chans, rescale=None, mask_threshold=-1, flow_threshold=0,\n",
    "                                                      omni=omni[2],cluster=cluster, resample=resample, tile=False, transparency=transparency, verbose=verbose)\n",
    "        # worm_bact_omni\n",
    "        masks[3], flows[3], styles[3] = model[2].eval(imglist, channels=chans, rescale=None, mask_threshold=-1, flow_threshold=0,\n",
    "                                                      omni=omni[3], cluster=cluster, resample=resample, tile=False, transparency=transparency, verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save0 = os.path.join('/home/kcutler/DataDrive/omnipose_all',dataset+'_comparison')\n",
    "io.check_dir(save0)\n",
    "print('Clean flag is {}'.format(clean))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if clean:\n",
    "    for j in range(n):\n",
    "        savedir = os.path.join(save0,names[j])\n",
    "        print(savedir)\n",
    "        io.check_dir(savedir)\n",
    "        io.save_masks(imgs, masks[j], flows[j], \n",
    "                      img_names,\n",
    "                      save_flows=True, \n",
    "                      save_outlines=0, \n",
    "                      savedir=savedir, \n",
    "                      in_folders=True, \n",
    "                      save_txt=False,\n",
    "                      save_plot=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in results\n",
    "This step also will also delete small masks from image boundaries. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from omnipose import utils \n",
    "\n",
    "for j in range(n):\n",
    "    savedir = os.path.join(save0,names[j])\n",
    "    mask_names = [os.path.join(savedir,'masks',os.path.splitext(os.path.basename(name))[0]+'_cp_masks.png') for name in img_names]\n",
    "    # flow_names = [os.path.join(savedir,'flows',os.path.splitext(os.path.basename(name))[0]+'_flows.tif') for name in img_names] #should fix png vs tif\n",
    "    masks[j] = [utils.format_labels(utils.clean_boundary(skimage.io.imread(f))).astype(np.uint32) for f in mask_names]\n",
    "    # flows[j] = [skimage.io.imread(f) for f in flow_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cellpose_omni import metrics\n",
    "from skimage import measure\n",
    "threshold=np.linspace(0.5,1,100)\n",
    "import fastremap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Note that 'ap' is called average precision in cellpose, but really it is the jaccard index\n",
    "ap = [[]] * len(masks) \n",
    "tp = [[]] * len(masks)\n",
    "fp = [[]] * len(masks)\n",
    "fn = [[]] * len(masks)\n",
    "IoU = [[]] * len(masks)\n",
    "OvR = [[]] * len(masks)\n",
    "pred_areas = [[]] * len(masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do the same boundary cleaning on ground truth masks \n",
    "import ncolor \n",
    "\n",
    "nimg = len(masks_gt)\n",
    "cell_areas = [[]] * nimg\n",
    "masks_gt_clean = [None]*nimg\n",
    "masks_pred_clean = [[None]*nimg]*len(masks)\n",
    "# remapping = [[]] * nimg\n",
    "for j in range(nimg):\n",
    "    mgt = ncolor.format_labels(utils.clean_boundary(masks_gt[j]))\n",
    "    masks_gt_clean[j] = mgt\n",
    "    regions = measure.regionprops(mgt)\n",
    "    areas = np.array([reg.area for reg in regions])\n",
    "    cell_areas[j] =  areas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.min(np.concatenate(cell_areas))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute performance metrics for each model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# go over each model\n",
    "if clean:\n",
    "    for j,masks_pred in enumerate(masks):\n",
    "    #     masks_pred = map(list,zip(*[ utils.format_labels(utils.clean_boundary(msk)) for msk in masks_pred]))\n",
    "    # just apply cleanup to the masks when reading them in \n",
    "        api,tpi,fpi,fni = metrics.average_precision(masks_gt_clean,masks_pred,threshold=threshold)\n",
    "        ap[j] = ap[j]+[api]\n",
    "        tp[j] = tp[j]+[tpi]\n",
    "        fp[j] = fp[j]+[fpi]\n",
    "        fn[j] = fn[j]+[fni]\n",
    "\n",
    "        masks_pred_clean[j] = masks_pred\n",
    "        # go over every image\n",
    "        for k in range(nimg):\n",
    "            # get the IoU matrix; axis 0 corresponds to GT, axis 1 to pred \n",
    "            regions = measure.regionprops(masks_pred[k])\n",
    "            areas = np.array([reg.area for reg in regions])\n",
    "            pred_areas[j] = pred_areas[j] + [areas]\n",
    "            iou = metrics._intersection_over_union(masks_gt_clean[k], masks_pred[k])\n",
    "            ovp = metrics._label_overlap(masks_gt_clean[k], masks_pred[k])[1:,1:] #throw out columns corresponding to zero  \n",
    "            OvR[j] = OvR[j]+[ovp / areas[np.newaxis,:]] # Overlap Ratio           \n",
    "            IoU[j] = IoU[j]+[iou]\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save0 = os.path.join('/home/kcutler/DataDrive/omnipose_all',dataset+'_comparison')\n",
    "io.check_dir(save0)\n",
    "\n",
    "if clean:\n",
    "    np.savez(os.path.join(save0,'OvR'),OvR)\n",
    "    np.savez(os.path.join(save0,'IoU'),IoU)\n",
    "    np.savez(os.path.join(save0,'cell_areas'),cell_areas)\n",
    "    np.savez(os.path.join(save0,'ap'),ap)\n",
    "# np.savez(savedir+'remapping',remapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read back in performance metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "OvR = np.load(os.path.join(save0,'OvR'+'.npz'),allow_pickle=True)['arr_0']\n",
    "IoU = np.load(os.path.join(save0,'IoU'+'.npz'),allow_pickle=True)['arr_0']\n",
    "cell_areas = np.load(os.path.join(save0,'cell_areas'+'.npz'),allow_pickle=True)['arr_0']\n",
    "ap = np.load(os.path.join(save0,'ap'+'.npz'),allow_pickle=True)['arr_0']\n",
    "\n",
    "if dataset=='bact_fluor':\n",
    "    savedir_phase = os.path.join('/home/kcutler/DataDrive/omnipose_all',datasets[0]+'_comparison')\n",
    "    ap_phase = np.load(os.path.join(savedir_phase,'ap'+'.npz'),allow_pickle=True)['arr_0']\n",
    "    ap_plot = [a for a in ap_phase]+[a for a in ap]\n",
    "    names_plot = ['phase_'+s for s in names_phase]+['fluor_'+s for s in names]\n",
    "\n",
    "else:\n",
    "    names_plot = names\n",
    "    ap_plot = [a for a in ap]\n",
    "    \n",
    "cell_count = np.array([len(np.unique(ca)) for ca in cell_areas])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(ap),j,n,names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot JI averaged over entire dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "from omnipose.utils import sinebow\n",
    "\n",
    "# computing the JI for the whole dataset by averaging per image\n",
    "# vs averagaing by cell area (I think the latter makes more sense, but it is not standard)\n",
    "# (also, it doesn't change the plot that much \n",
    "weighted = 0\n",
    "\n",
    "x = threshold\n",
    "sz = 1\n",
    "golden = (1 + 5 ** 0.5) / 2\n",
    "labelsize = 7\n",
    "\n",
    "\n",
    "# colors = ['g','r','b','y','c','m']\n",
    "n = len(names_plot)\n",
    "z = 1\n",
    "master_color_scheme = [[i,0,0] for i in np.linspace(1,.5,z)]+[[i,i,i] for i in np.linspace(.75,.25,n-z)]\n",
    "\n",
    "darkmode = 0\n",
    "if darkmode:\n",
    "    plt.style.use('dark_background')\n",
    "    axcol = 'w'\n",
    "    colors = sinebow(n+1)\n",
    "    colors = [colors[j+1] for j in range(n)]\n",
    "    background_color = 'k'\n",
    "    suffix = '_dark'\n",
    "else:\n",
    "    mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "    axcol = 'k'\n",
    "    colors = master_color_scheme\n",
    "    # if weighted:\n",
    "    #     colors = sinebow(n+1)\n",
    "    #     colors = [colors[j+1] for j in range(n)]\n",
    "    background_color = np.array([1,1,1,1])\n",
    "    suffix = ''\n",
    "\n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "fig = plt.figure(figsize=(sz, sz)) \n",
    "ax = plt.axes()\n",
    "# plt.tight_layout()\n",
    "# plt.minorticks_on()\n",
    "plt.xticks(np.arange(min(x), max(x)+1, .25))\n",
    "plt.xlim([0.5,1])\n",
    "plt.ylim([0,1])\n",
    "# plt.yticks(np.arange(0, 1.1, .25))\n",
    "\n",
    "selection = range(n)\n",
    "print(names_plot)\n",
    "if dataset=='bact_fluor':\n",
    "    # selection = [-3,-2,-1]\n",
    "    selection = range(len(ap_plot))\n",
    "    # selection = [-1,-2]\n",
    "for j in selection:\n",
    "    \n",
    "    if weighted:\n",
    "        mJI = np.sum(np.array([ji*np.sum(a) for a,ji in zip(cell_areas,ap_plot[j][0])]),axis=0)/np.sum(np.concatenate(cell_areas))\n",
    "        ax.plot(x,mJI.T,'--',label=names_plot[j],color=colors[j],linewidth=.5)\n",
    "\n",
    "    ax.plot(x,np.mean(ap_plot[j][0],axis=0).T,label=names_plot[j],color=colors[j],linewidth=.5)\n",
    "#     ax.plot(x,np.mean((fp[j][0]+fn[j][0]),axis=0),label=pretty_names[j],color=colors[j])\n",
    "\n",
    "\n",
    "# ax.set_facecolor('w')\n",
    "ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.05, 1))\n",
    "# ax.legend(prop={'size': labelsize}, loc='best', frameon=False)\n",
    "ax.tick_params(axis='both', which='major', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "ax.tick_params(axis='both', which='minor', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "# plt.setp(ax.xaxis.get_label(), visible=True, text='IoU')\n",
    "# plt.setp(ax.get_xticklabels(), visible=True, ha='right')\n",
    "\n",
    "# ax.grid(b=True, which='major', color='b', linestyle='-')\n",
    "ax.set_ylabel('Jaccard Index', fontsize = labelsize)\n",
    "ax.set_xlabel('IoU matching threshold', fontsize = labelsize)\n",
    "# plt.set(xlabel='IoU threshold', ylabel='Average Precision',fontsize=labelsize)\n",
    "# plt.tight_layout()\n",
    "# plt.yscale('log')\n",
    "\n",
    "\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.patch.set_alpha(0.0)\n",
    "\n",
    "fig.patch.set_facecolor(background_color)\n",
    "plt.show()\n",
    "\n",
    "figname = 'JI_vs_IoU'\n",
    "# fig.savefig(os.path.join(save0,figname+suffix+'.eps'),bbox_inches=tight_bbox)\n",
    "# fig.savefig(os.path.join(save0,figname+suffix+'.pdf'),bbox_inches=tight_bbox)\n",
    "for ext in ['.pdf','.png','.eps']:\n",
    "    fig.savefig(os.path.join(save0,figname+suffix+ext),bbox_inches=\"tight\",pad_inches = 0.05)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Calculate errors "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = len(names)\n",
    "nimg = len(cell_areas)\n",
    "cell_errors = [[]]*N\n",
    "total_errors = [0]*N\n",
    "total_single_errors = [0]*N\n",
    "total_cells = len(np.concatenate(cell_areas))\n",
    "\n",
    "# area_thresh = np.linspace(1,np.max(np.concatenate(cell_areas)),100)\n",
    "# area thresholds \n",
    "area_thresh = [np.percentile(np.concatenate(cell_areas),i) for i in range(100)]\n",
    "M = len(area_thresh)\n",
    "percent = [0]*N\n",
    "\n",
    "\n",
    "for j in range(N):\n",
    "    print(names[j])\n",
    "    ce_thresh = [[]*2]*M\n",
    "    te_thresh = np.zeros((2,M))\n",
    "    tse_thresh = np.zeros((2,M))\n",
    "    tc_thresh = np.zeros((2,M))\n",
    "    \n",
    "    for k in range(nimg):\n",
    "        r = OvR[j][k].copy() #overlap ratio \n",
    "        r[r<=0.75] = 0 \n",
    "        mx = np.max(r,axis=0) #find the maximum overlap for each true cell \n",
    "        mx[mx==0] = np.nan #exclude case where a spurious label has no overlaps\n",
    "        hits = np.sum(r==mx,axis=1) # sum will be zero if a cell label has zero hits, not sure which ones these are\n",
    "        cell_errors[j] = cell_errors[j] + [np.abs(hits-1)] #error if >1 or =0 \n",
    "        total_errors[j] += np.sum(hits[hits>1]-1)+np.sum(hits==0) # -1 because a hit of 2 is 2 pred labels, 1 'extra' = 1 error\n",
    "        total_single_errors[j] += np.sum(hits>1)+np.sum(hits==0)\n",
    "        \n",
    "        ca = cell_areas[k]\n",
    "        for i,a in enumerate(area_thresh):\n",
    "           \n",
    "            for l in range(2): # look at both above AND below the area threshold\n",
    "                cell_filter = ca>=a if l==0 else ca<a\n",
    "                cell_count = np.count_nonzero(cell_filter)\n",
    "                hits_thresh = hits[cell_filter] # sum will be zero if a cell label has zero hits, not sure which ones these are\n",
    "#                 ce_thresh[l,i] += [np.abs(hits_thresh-1)] #error if >1 or =0 \n",
    "                te_thresh[l,i] += np.sum(hits_thresh[hits_thresh>1]-1)+np.sum(hits_thresh==0) # -1 because a hit of 2 is 2 pred labels, 1 'extra' or 1 error\n",
    "                tse_thresh[l,i] += np.sum(hits_thresh>1)+np.sum(hits_thresh==0)\n",
    "                tc_thresh[l,i] += cell_count\n",
    "\n",
    "    \n",
    "    eps = 1e-8\n",
    "    percent[j] = [[utils.safe_divide(tse_thresh[l,i],tc_thresh[l,i]+eps)*100 for i in range(M)] for l in range(2)]        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Print and save error metrics  \n",
    "\n",
    "Calculate probability of errors above and below cutoff\n",
    "This is used in the section \"XXX\". We save the output to a test file in the `save0` directory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = np.concatenate(cell_areas)\n",
    "m = np.percentile(x,75)\n",
    "top_count = np.count_nonzero(x>=m)\n",
    "bottom_count = np.count_nonzero(x<m)\n",
    "cell_count = np.count_nonzero(x)\n",
    "string = ['Top quartile of area is {} square pixels'.format(m),\n",
    "         'Number of cells in top quartile is {}'.format(top_count),\n",
    "         'Number of cells in bottom three quartiles is {}'.format(bottom_count),\n",
    "          '']\n",
    "# for j in [2,1,0]:\n",
    "for j in range(N):\n",
    "    y = np.concatenate(cell_errors[j])\n",
    "    p = y>0\n",
    "    q = np.logical_and(x>=m,p)\n",
    "    s = np.logical_and(x<m,p)\n",
    "    N_top = np.count_nonzero(q)\n",
    "    N_bottom = np.count_nonzero(s)\n",
    "    N = np.count_nonzero(y)\n",
    "\n",
    "    substr = [names[j],\n",
    "            'Number/fraction of cells in the top quartile with one or more errors: {} / {}'.format(N_top,N_top/top_count*100),\n",
    "            'Number/fraction of cells in the bottom three quartiles with one or more errors: {} / {}'.format(N_bottom, N_bottom/bottom_count*100),\n",
    "            'Number/fraction of all cells with one or more errors: {} / {} \\n'.format(N, N/cell_count*100)]\n",
    "    string += substr\n",
    "    \n",
    "# see the output in the notebook and save to txt file\n",
    "with open(os.path.join(save0,'error_metrics.txt'), \"w\") as text_file:\n",
    "    for s in string:\n",
    "        [print(s,file=f) for f in [text_file,None]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Plot Figure 4A\n",
    "Plot percent of cells with >=1 error against the area percentile cutoff. This tells us if there are any differences in how the algorithm treats cells by area. A flat line indicates that there is no difference in error count between large and small cells. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "axcol = 'k'\n",
    "\n",
    "from omnipose.utils import sinebow\n",
    "n = len(names)\n",
    "z = 1\n",
    "master_color_scheme = [[i,0,0] for i in np.linspace(1,.5,z)]+[[i,i,i] for i in np.linspace(.75,0,n-z)]\n",
    "\n",
    "golden = (1 + 5 ** 0.5) / 2\n",
    "sz = 1.5\n",
    "labelsize = 7\n",
    "\n",
    "%matplotlib inline\n",
    "darkmode = 0\n",
    "if darkmode:\n",
    "    plt.style.use('dark_background')\n",
    "    axcol = 'w'\n",
    "    colors = sinebow(n+1)\n",
    "    colors = [colors[j+1] for j in range(n)]\n",
    "    background_color = 'k'\n",
    "    suffix = '_dark'\n",
    "else:\n",
    "    mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "    axcol = 'k'\n",
    "    cmap = mpl.cm.get_cmap('viridis')\n",
    "#     colors = cmap(np.linspace(0,.9,len(names)))\n",
    "    colors = master_color_scheme\n",
    "    background_color = np.array([1,1,1,1])\n",
    "    suffix = ''\n",
    "    \n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "\n",
    "# could turn these results into a plot by repeating at different cutoffs \n",
    "fig = plt.figure(figsize=(sz, sz)) \n",
    "ax = plt.axes()\n",
    "for j in range(n):\n",
    "# for j in [-1,-2]:\n",
    "# for j in [0,1,-1]:\n",
    "#     ax.plot(area_thresh,np.divide(tse_thresh[j],tc_thresh),label=names[j],color=colors[j])\n",
    "#     ax.plot(range(100),percent[j],label=names[j],color=colors[j])\n",
    "    ax.plot(range(100),percent[j][0],label=names[j],color=colors[j])\n",
    "#     ax.plot(range(100),percent[j][1],'--',label='above',color=colors[j])\n",
    "#     print('mean percent error is',np.mean(percent[j]))\n",
    "    # ax.set_facecolor('w')\n",
    "# plt.xscale('log')\n",
    "# ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.05, 1))\n",
    "\n",
    "# ax.legend(prop={'size': labelsize}, loc='best', frameon=False)\n",
    "ax.tick_params(axis='both', which='major', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "ax.tick_params(axis='both', which='minor', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "ax.set_ylabel('Percent of cells \\nwith >=1 error', fontsize = labelsize)\n",
    "ax.set_xlabel('Area percentile cutoff', fontsize = labelsize)\n",
    "plt.ylim([0,100])\n",
    "plt.xlim([0,100])\n",
    "\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.patch.set_alpha(0.0)\n",
    "ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.05, 1))\n",
    "fig.patch.set_facecolor(background_color)\n",
    "plt.show()\n",
    "\n",
    "a = 35\n",
    "tight_bbox_raw = ax.get_tightbbox(fig.canvas.get_renderer())\n",
    "tight_bbox_raw._points+=[[-a,-a],[a,a]]\n",
    "tight_bbox = mpl.transforms.TransformedBbox(tight_bbox_raw, mpl.transforms.Affine2D().scale(1./fig.dpi))\n",
    "\n",
    "\n",
    "\n",
    "figname = 'Error_vs_area'\n",
    "# fig.savefig(os.path.join(save0,figname+suffix+'.eps'),bbox_inches=tight_bbox)\n",
    "# fig.savefig(os.path.join(save0,figname+suffix+'.pdf'),bbox_inches=tight_bbox)\n",
    "for ext in ['.pdf','.png','.eps']:\n",
    "    fig.savefig(os.path.join(save0,figname+suffix+ext),bbox_inches=\"tight\",pad_inches = 0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Consolidate error and cell area for plotting \n",
    "a = np.concatenate(cell_areas)\n",
    "\n",
    "area_cutoff = np.percentile(a,75)\n",
    "cell_filter = a>=area_cutoff\n",
    "text = ['number of cells: {}'.format(len(a)),\n",
    "        'fraction of areas: {}'.format(np.count_nonzero(cell_filter)/len(a)*100),\n",
    "        'top area quartile: {}\\n'.format(area_cutoff)]\n",
    "# for j in [2,1,0]:\n",
    "\n",
    "for j in range(n):\n",
    "    e = np.concatenate(cell_errors[j]) \n",
    "    p = np.array(percent[j][0])\n",
    "    text+=[names[j],'fraction of errors: {}'.format(np.sum(e[cell_filter])/np.sum(e)*100),\n",
    "           'minimum error percentage: {}'.format(np.min(p[p>0])),\n",
    "           'maximum error percentage: {}'.format(np.max(p[p>0])),\n",
    "           # 'cross 5 percent: {}'.format(np.argwhere(p>5)),\n",
    "           'mean error percentage: {}\\n'.format(np.mean(p[p>0]))]\n",
    "\n",
    "with open(os.path.join(save0,'error_fraction_by_area.txt'), \"w\") as text_file:\n",
    "    [print(t,file=f) for t in text for f in [text_file,None]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Separate into cell types for plotting statistics\n",
    "\n",
    "This section of the notebook is only intended to be run for the bacterial datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.scatter(IoU[1])\n",
    "exclude = 'xxxx' # could explude some problematic images for all \n",
    "\n",
    "if dataset=='bact_phase':\n",
    "    cats = ['regular morphologies','mutants, antibiotics','elongated wildtype']\n",
    "    subsets= [['5I_crop','PAO1_Staph','PSVB','Serratia_Ecoli','wiggins','vibrio','bthai'],\n",
    "              ['Hpylori','dnaA','ftsN','murA','cex','a22','Az'],['caulo','streptomyces']]\n",
    "\n",
    "elif dataset=='bact_fluor':\n",
    "    cats = ['regular morphologies','mutants and antibiotics']\n",
    "    subsets= [['wiggins','vibrio','bthai'],\n",
    "              ['cex','a22']]\n",
    "elif dataset=='worm':\n",
    "    cats = ['worms']\n",
    "    subsets= [['elegans']]\n",
    "else:\n",
    "    cats = ['']\n",
    "    subsets = [['']]\n",
    "K = len(cats)\n",
    "indices = [np.array([i for i,s in enumerate(mask_names) if any((name in s) and (exclude not in s) for name in subsets[k])]).astype(int) for k in range(K)]\n",
    "\n",
    "counts = np.zeros((K,1),int)\n",
    "for j,inds in enumerate(indices):\n",
    "    # print(j)\n",
    "    for i in inds:\n",
    "        # print(len(np.unique(masks_gt_clean[i]))-1)\n",
    "        counts[j] +=len(np.unique(masks_gt_clean[i]))-1\n",
    "counts\n",
    "\n",
    "\n",
    "print('Are all images accounted for?',np.sum([len(indices[k]) for k in range(K)])==len(mask_names))\n",
    "print('Image category counts:',[len(indices[k]) for k in range(K)])\n",
    "print('Image category cell counts:',counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# len(mask_names),np.sum([len(indices[k]) for k in range(K)])\n",
    "# indices\n",
    "# mask_names[0]\n",
    "np.sum([len(indices[k]) for k in range(K)]),len(mask_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Jaccard Index averages "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "per_cell = 1\n",
    "J = len(names)\n",
    "K = len(cats)\n",
    "# y = [[]*K]*J #BADDDDDDDDDD! Need to check my other results to make sure it was done correctly, https://stackoverflow.com/questions/54673821/python-how-to-initialize-a-nested-list-with-empty-values-which-i-can-append-to \n",
    "y = [ [ [] for i in range(K) ] for i in range(J) ]\n",
    "for j in range(J):\n",
    "    for k in range(K):\n",
    "        matched_iou = []\n",
    "        mean_matched_iou = []\n",
    "        for ind in indices[k]:\n",
    "            iou = IoU[j][ind][1:,1:]\n",
    "            th = 0\n",
    "            n_min = min(iou.shape[0], iou.shape[1])\n",
    "            costs = -(iou >= th).astype(float) - iou / (2*n_min)\n",
    "            true_ind, pred_ind = linear_sum_assignment(costs)\n",
    "            miou = iou[true_ind, pred_ind]\n",
    "            matched_iou.append(miou)\n",
    "            mean_matched_iou.append(np.sum(miou)/max(iou.shape[0], iou.shape[1]))\n",
    "            \n",
    "        y[j][k] = [m for sublist in matched_iou for m in sublist] if per_cell else mean_matched_iou"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Jaccard Index per image\n",
    "Color-coded by cell type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "from omnipose.utils import sinebow\n",
    "J = len(names)\n",
    "linestyle='-'\n",
    "\n",
    "%matplotlib inline\n",
    "darkmode = 0\n",
    "if darkmode:\n",
    "    plt.style.use('dark_background')\n",
    "    axcol = 'w'\n",
    "    colors = sinebow(J+1)\n",
    "    colors = [colors[j+1] for j in range(n)]\n",
    "    background_color = 'k'\n",
    "else:\n",
    "    mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "    axcol = 'k'\n",
    "    cmap = mpl.cm.get_cmap('viridis')\n",
    "    colors = cmap(np.linspace(0,.9,len(names)))\n",
    "#     colors = master_color_scheme\n",
    "    background_color = np.array([1,1,1,1])\n",
    "    \n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "\n",
    "x = threshold\n",
    "golden = (1 + 5 ** 0.5) / 2\n",
    "sz = 6.5\n",
    "labelsize = 7\n",
    "# fig = plt.figure(figsize=(sz, sz/golden)) \n",
    "# fig = plt.figure(figsize=(sz, sz/2)) \n",
    "# ax = plt.axes()\n",
    "fig, axs = plt.subplots(1,len(names),figsize=(sz,sz/J),sharex=True, sharey=True)\n",
    "\n",
    "# colors = ['tab:olive','tab:gray','tab:red']\n",
    "# plt.tight_layout()\n",
    "# plt.minorticks_on()\n",
    "plt.xticks(np.arange(min(x), max(x)+1, .25))\n",
    "# plt.xlim([0.45,1.05])\n",
    "# plt.ylim([0,1])\n",
    "# plt.yticks(np.arange(0, 1.1, .25))\n",
    "\n",
    "# colors = ['g','r','b','y','c','m']\n",
    "# pretty_names = ['Original Cellpose','','','Mixed Method','','','','New Method']\n",
    "cell_count = np.array([len(np.unique(ca)) for ca in cell_areas])\n",
    "# for j in range(n):\n",
    "# sort_colors = [colors[j] for j in args]\n",
    "# colors = ['r',[.75,.75,.75],[0,0,0],[.5,.5,.5],[1,0,0]]\n",
    "\n",
    "fill = 0 # makes mean+-\"err\" fill around the average JI\n",
    "per_image = 1 # plots JI per image\n",
    "density = 0 # plots JI per image as a dot for each IoU and colors by density \n",
    "\n",
    "l = np.array([len(indices[k]) for k in range(K)])\n",
    "# alpha = 1.5/(1+l/np.min(l))\n",
    "alpha = [1,1,1]\n",
    "\n",
    "from scipy.stats import gaussian_kde\n",
    "\n",
    "for j in range(J):\n",
    "    \n",
    "    mean = np.mean(ap[j][0],axis=0).T\n",
    "    err = np.std(ap[j][0],axis=0).T\n",
    "    if len(names)>1:\n",
    "        ax = axs[j]\n",
    "    else:\n",
    "        ax = axs\n",
    "    if fill:\n",
    "        ax.plot(x,mean,label=names[j],color=colors[j],linestyle=linestyle)\n",
    "        ax.fill_between(x,mean-err,mean+err,facecolor=colors[j],alpha=0.5)\n",
    "    if per_image:\n",
    "#         ax.plot(x,ap[j][0].T,label=names[j],color=colors[j],linestyle=linestyle,alpha=.05)\n",
    "#         ax.plot(x,mean,label=names[j],color='k',linestyle=linestyle)\n",
    "        for k in range(K):\n",
    "#             if k<K-1:\n",
    "#                 alpha = .25\n",
    "#             else:\n",
    "#                 alpha = .5\n",
    "            arr = ap[j][0][indices[k]].T\n",
    "            ax.plot(x,arr,label=names[j],color=colors[j],linestyle=linestyle,alpha=alpha[k],lw=.75)\n",
    "            ax.plot(x,mean,label=names[j],color=axcol,linestyle=linestyle,lw=.75)\n",
    "    if density:\n",
    "        y = np.array(ap[j]).T.flatten()\n",
    "        X = np.repeat(x,len(ap[j][0]))\n",
    "        xy = np.vstack([X,y])\n",
    "        z = gaussian_kde(xy)(xy)\n",
    "        idx = z.argsort()\n",
    "        X, y, z = X[idx], y[idx], z[idx]\n",
    "        ax.scatter(X, y, c=np.log(z), s=1)\n",
    "    # ax.set_facecolor('w')\n",
    "    # ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.05, 1))\n",
    "#     ax.legend(prop={'size': 5}, loc='best', frameon=False)\n",
    "    ax.set_title(names[j],fontsize=labelsize)\n",
    "    ax.tick_params(axis='both', which='major', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "    ax.tick_params(axis='both', which='minor', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True)\n",
    "#     ax.tick_params(axis='x', which='both', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=0,left=0,labelbottom=0)\n",
    "\n",
    "#     if j==J-1:\n",
    "#         ax.set_xlabel('IoU threshold', fontsize = labelsize)\n",
    "#         ax.tick_params(axis='x', which='both', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True,labelbottom=1)\n",
    "#         ax.tick_params(axis='x', which='minor', labelsize=labelsize,length=3, direction=\"out\",colors=axcol,bottom=True,left=True,labelbottom=1)\n",
    "\n",
    "#     ax.set_ylim((0,1))\n",
    "    ax.set_xlim((.5-.05,1))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.patch.set_alpha(0.0)\n",
    "\n",
    "plt.subplots_adjust(bottom=.3,left=.1)\n",
    "fig.patch.set_facecolor(background_color)\n",
    "fig.supylabel('Jaccard Index', fontsize = labelsize)\n",
    "fig.supxlabel('IoU matching threshold', fontsize = labelsize)\n",
    "# swap last two \n",
    "# pos = axs[-1].get_position()\n",
    "# axs[-1].set_position(axs[-2].get_position())\n",
    "# axs[-2].set_position(pos)\n",
    "# \n",
    "# plt.set(xlabel='IoU threshold', ylabel='Average Precision',fontsize=labelsize)\n",
    "# plt.tight_layout()\n",
    "# plt.yscale('log')\n",
    "\n",
    "\n",
    "\n",
    "plt.show()\n",
    "\n",
    "tight_bbox_raw = ax.get_tightbbox(fig.canvas.get_renderer())\n",
    "a = 50\n",
    "tight_bbox_raw._points+=[[-a,-a],[a,a]]\n",
    "tight_bbox = mpl.transforms.TransformedBbox(tight_bbox_raw, mpl.transforms.Affine2D().scale(1./fig.dpi))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute fractional difference between Omnipose and Cellpose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cutoff = 80 # IoU cutoff of 0.8\n",
    "if dataset=='bact_fluor':\n",
    "    i1 = 0\n",
    "    i2 = 1\n",
    "else:\n",
    "    i1 = 0\n",
    "    i2 = 2\n",
    "print('Comparing models {} and {}'.format(names[i1],names[i2]))\n",
    "cp = np.mean(ap[i1][0],axis=0).T\n",
    "om = np.mean(ap[i2][0],axis=0).T\n",
    "\n",
    "av_all = np.mean(np.divide((om-cp),cp,out=np.zeros_like(cp), where=np.logical_and(cp!=0,~np.isnan(cp))))\n",
    "av_up = np.mean(np.divide((om[cutoff:]-cp[cutoff:]),cp[cutoff:],out=np.zeros_like(cp[cutoff:]), where=np.logical_and(cp[cutoff:]!=0,~np.isnan(cp[cutoff:]))))\n",
    "av_dwn = np.mean(np.divide((om[:cutoff]-cp[:cutoff]),cp[:cutoff],out=np.zeros_like(cp[:cutoff]), where=np.logical_and(cp[:cutoff]!=0,~np.isnan(cp[:cutoff]))))\n",
    "\n",
    "text = ['average percent difference across all IoUs: {}'.format(av_all*100),\n",
    "        'percent difference above and below IoU={}: {},{}'.format(cutoff/100,av_up*100,av_dwn*100),\n",
    "        'percent difference of averages across all IoUs: {}'.format((np.mean(om)-np.mean(cp))/np.mean(cp)*100)] # turn this into saving a text file \n",
    "\n",
    "with open(os.path.join(save0,'relative_performance.txt'), \"w\") as text_file:\n",
    "    [print(t,file=f) for t in text for f in [text_file,None]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Figure 4f - worm dataset violin plot\n",
    "To run this section, ensure the worm dataset has been loaded and all non-dataset-specific cells above have been run. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sz = 2\n",
    "lw = 0.8\n",
    "fig, axs = plt.subplots(figsize=(sz,sz))\n",
    "plt.subplots_adjust(wspace=.4, hspace=.4)\n",
    "# axs = axs.flatten()\n",
    "np.random.seed(123)\n",
    "w0 = 0.7\n",
    "split=1\n",
    "box = False\n",
    "scatter = 0\n",
    "violin = 1\n",
    "\n",
    "\n",
    "# for j in [-1]:\n",
    "ax = axs\n",
    "\n",
    "# x = np.linspace(0,1,len(y))\n",
    "text = []\n",
    "for j in range(J):\n",
    "    percents = [np.sum(np.array(yi)>0.8)/counts[k]*100 for k,yi in enumerate(y[j])]\n",
    "    overall_percent = np.sum([np.sum(np.array(yi)>0.8) for yi in y[j]]) / np.sum(counts)\n",
    "    # text+=[names[j], 'Percent above 0.8',percents,overall_percent]#,height,len(y[j]))\n",
    "    text += [names[j],'Percent above 0.8:\\nBy category: {} \\nOverall: {} \\n'.format(percents,overall_percent)]#,height,len(y[j]))\n",
    "\n",
    "w = w0 / 3\n",
    "j = 0\n",
    "x = np.linspace(j-w,j+w,len(y))\n",
    "height = [np.mean(yi) for yi in y[j]]\n",
    "\n",
    "yp = [np.array(yi).flatten() for yi in y]\n",
    "\n",
    "parts = ax.violinplot(yp,positions=x,showextrema=0,widths=w0/(len(y)+1)) #showmeans=1,quantiles = [[.25,.75]]*len(x),\n",
    "boxes = ax.boxplot(yp,positions=x,showfliers=False,widths=w0/35,patch_artist=True,\n",
    "                   boxprops=dict(facecolor='w', color='k',linewidth=lw),whiskerprops=dict(color='k',linewidth=lw),\n",
    "                   medianprops=dict(color='k',linewidth=lw))\n",
    "\n",
    "for i,pc in enumerate(parts['bodies']):\n",
    "    pc.set_facecolor(colors[i])\n",
    "    pc.set_alpha(1)\n",
    "    pc.set_edgecolor(colors[i])\n",
    "    pc.set_linewidth(lw)\n",
    "\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.patch.set_alpha(0.0)\n",
    "\n",
    "ax.set_ylim(0,1)      \n",
    "ax.set_xlim((j-2*w0/4,j+2*w0/4))\n",
    "ax.set_xticks([])\n",
    "plt.show()   \n",
    "tight_bbox_raw = ax.get_tightbbox(fig.canvas.get_renderer())\n",
    "a = 50\n",
    "tight_bbox_raw._points+=[[-a,-a],[a,a]]\n",
    "tight_bbox = mpl.transforms.TransformedBbox(tight_bbox_raw, mpl.transforms.Affine2D().scale(1./fig.dpi))\n",
    "\n",
    "fig.savefig(os.path.join(save0,'Iou_violin.pdf'),bbox_inches=tight_bbox)\n",
    "\n",
    "with open(os.path.join(save0,'IoU_stats.txt'), \"w\") as text_file:\n",
    "    [print(t,file=f) for t in text for f in [None,text_file]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Segmentation error examples\n",
    "Most of what we observe in small cells are ambiguous division events - cells that may or may not be divided. 1 error is either a missing mask / 2 masks merged into one (not yet divided even though GT says it is) or split into 2 masks (has divided even though GT says it has not). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ncolor \n",
    "from cellpose_omni import transforms\n",
    "import cellpose, omnipose \n",
    "\n",
    "#code developing the error ideas \n",
    "# NOTE: broken for flows that are not rescaled, need to add zoom command\n",
    "ind = range(4)\n",
    "\n",
    "from matplotlib.colors import ListedColormap\n",
    "plt.style.use('dark_background')\n",
    "# now find and crop specific cells for comparison \n",
    "# oversegmneted cells will have low IoUs spread out over several hits; let's take a specific image from the rpevious list as an example. We can see that there are three true labels, but \n",
    "# several predicted labels. The IoU can be at most 1, but small precited labels will always be quite small, up to 1/(area of the true label). But we don't really want to care about the \n",
    "cm2 = ListedColormap([color for color in sinebow(5).values()])\n",
    "\n",
    "k = 0\n",
    "j = -1 # look at Omnipose\n",
    "\n",
    "m = 0\n",
    "r_pred = OvR[j][k].copy()\n",
    "mx = np.max(r_pred,axis=0) \n",
    "# std = np.std(r_pred,axis=0)\n",
    "mx[mx==0] = np.nan\n",
    "hits_pred = r_pred==mx\n",
    "err = np.abs(np.sum(hits_pred,axis=1)-1)\n",
    "\n",
    "# err = cell_errors[j][k]# assuming proper precomputing \n",
    "areas = cell_areas[k]\n",
    "error_cutoff = 1\n",
    "area_cutoff = 10\n",
    "err_indexes = np.nonzero(np.logical_and(err>=error_cutoff,areas>area_cutoff))[0]\n",
    "\n",
    "pad = 5\n",
    "\n",
    "for l in err_indexes[0:10]:\n",
    "    print(err[l])\n",
    "    mgt = masks_gt_clean[k]\n",
    "    y,x = np.nonzero(mgt==l+1)\n",
    "    max_y,max_x = np.array(mgt.shape)-1\n",
    "    \n",
    "    y0 = max(0,min(y)-pad)\n",
    "    y1 = min(max_y,max(y)+pad)\n",
    "    x0 = max(0,min(x)-pad)\n",
    "    x1 = min(max_x,max(x)+pad)\n",
    "    slc = (Ellipsis,)+omnipose.utils.bbox_to_slice([y0,y1,x0,x1],mgt.shape)\n",
    "    p = transforms.normalize99(imgs[k][slc],omni=True)\n",
    "    p = transforms.move_min_dim(p)\n",
    "    if p.ndim>2:\n",
    "        p=p[...,0]\n",
    "    p = p\n",
    "    \n",
    "    mask_gt = mgt[slc]\n",
    "    bini = mask_gt==l+1\n",
    "    img0 = np.stack([p]*3,axis=-1)\n",
    "    # outli = transforms.normalize99(plot.outline_view(img0*255,bini),omni=True)\n",
    "    bin0 = np.stack([bini]*3,axis=-1)\n",
    "    # gt_pic = np.hstack((img0,outli,bin0))\n",
    "    mask_pred = masks[j][k][slc]\n",
    "    # flow_pred = transforms.normalize99(flows[j][k][0][y0:y1,x0:x1],omni=True)[:,:,:3] #transparency an issue here, at least with concatenating \n",
    "    mu = flows[j][k][1][slc]\n",
    "    flow_pred = plot.dx_to_circ(mu)/255\n",
    "    \n",
    "    inds = np.where(hits_pred[l,:])\n",
    "    rmask = np.zeros_like(mask_gt)\n",
    "    for i,label in enumerate(inds[0]):\n",
    "        rmask[mask_pred==label+1] = i+1\n",
    "\n",
    "    \n",
    "    outl_pred = plot.outline_view(img0,rmask)/255\n",
    "    res_pred = np.hstack((img0,bin0,cm2(ncolor.label(mask_pred,max_depth=10))[:,:,:3],outl_pred,flow_pred))\n",
    "    plt.figure(figsize=[2,2])\n",
    "    plt.imshow(res_pred)\n",
    "    plt.axis('off')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extended Data Figure 4a\n",
    "For this plot, I should put both the training and test errors. The training dataset has a lot more long cells in it. In fact, I should do the analysis with the train, test, and combined to show correlation and lack of overfitting. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "import matplotlib.patches as patches\n",
    "from scipy.stats import gaussian_kde\n",
    "import cmasher as cmr\n",
    "import ncolor\n",
    "\n",
    "\n",
    "def custom_ceil(x,base=10):\n",
    "    return int(np.ceil(x / float(base)) * base)\n",
    "\n",
    "color = [1,0,0,.5]\n",
    "axcol = 'k'\n",
    "\n",
    "# some random data\n",
    "# x = np.random.randn(1000)\n",
    "# y = np.random.randn(1000)\n",
    "j = 0 # which algorithm\n",
    "x = np.concatenate(cell_areas)\n",
    "y = np.concatenate(cell_errors[j])\n",
    "xmin = 10\n",
    "xmax = 10**5\n",
    "# cm = plt.cm.get_cmap('plasma')\n",
    "cm = cmr.get_sub_cmap('plasma', 0.2, 0.95)\n",
    "\n",
    "mpl.rcParams['figure.dpi'] = 1000\n",
    "rasterized = True\n",
    "\n",
    "def scatter_hist(x, y, ax, ax_histx, ax_histy,zorder=1, hist = 1):\n",
    "    # no labels\n",
    "    ax_histx.tick_params(axis=\"x\", labelbottom=False)\n",
    "    ax_histy.tick_params(axis=\"y\", labelleft=False)\n",
    "\n",
    "    # the scatter plot:\n",
    "#     ax.scatter(x, y, s=3,edgecolors='none',c=color)\n",
    "    \n",
    "#         y = np.array(ap[j]).T.flatten()\n",
    "#         X = np.repeat(x,len(ap[j][0]))\n",
    "    X = x.copy()\n",
    "    xy = np.vstack([X,y])\n",
    "    z = gaussian_kde(xy)(xy)\n",
    "    idx = z.argsort()\n",
    "    X, y, z = X[idx], y[idx], z[idx]\n",
    "    ax.scatter(X, y, c=np.log(z), s=1,cmap=cm,zorder=zorder,rasterized=rasterized)\n",
    "    \n",
    "    \n",
    "    if hist:\n",
    "        # now determine limits by hand:\n",
    "        binwidth1 = 100\n",
    "        binwidth2 = 1\n",
    "        eps = 1e-20\n",
    "        xymax = max(np.max(np.abs(x)), np.max(np.abs(y)))\n",
    "        lim = (int(xymax/binwidth1) + 1) * binwidth1\n",
    "        pp = 1\n",
    "        bins1 = np.arange(xmin, lim + binwidth1, binwidth1)\n",
    "        bins2 = np.arange(0, 55 + binwidth2, binwidth2)\n",
    "        logbins1 = np.logspace(np.log10(bins1[0]),np.log10(bins1[-1]),len(bins1))\n",
    "\n",
    "        # plot histogram\n",
    "        n,bins,patches = ax_histx.hist(x, bins=logbins1,color=color)\n",
    "        col = (n-n.min())/(n.max()-n.min())+eps\n",
    "        for c, p in zip(col, patches):\n",
    "            plt.setp(p, 'facecolor', cm(c**pp))\n",
    "\n",
    "        n,bins,patches = ax_histy.hist(y, bins=bins2, orientation='horizontal',color=color,align='left')\n",
    "        col = (n-n.min())/(n.max()-n.min())+eps\n",
    "        for c, p in zip(col, patches):\n",
    "            plt.setp(p, 'facecolor', cm(c**.25))\n",
    "    \n",
    "    \n",
    "#     ax_histx.tick_params(axis='x',which='both',bottom=False,top=False)\n",
    "#     ax_histy.tick_params(axis='y',which='both',left=False,right=False)\n",
    "#     ax_histy.get_xaxis().set_ticks([])\n",
    "#     ax_histy.get_yaxis().set_ticks([])\n",
    "    \n",
    "\n",
    "#     plt.xlim([0,10**3])\n",
    "# start with a square Figure\n",
    "# a = 1+golden\n",
    "a = 5\n",
    "b = 1\n",
    "c = 2.5\n",
    "labelsize = 9\n",
    "fig = plt.figure(figsize=(c, c))\n",
    "xlog = 1\n",
    "\n",
    "# Add a gridspec with two rows and two columns and a ratio of 2 to 7 between\n",
    "# the size of the marginal axes and the main axes in both directions.\n",
    "# Also adjust the subplot parameters for a square plot.\n",
    "\n",
    "gs = fig.add_gridspec(2, 2,  width_ratios=(a, b), height_ratios=(b, a),\n",
    "                      left=0.1, right=0.9, bottom=0.1, top=0.9,\n",
    "                      wspace=0.075, hspace=0.075)\n",
    "\n",
    "ax = fig.add_subplot(gs[1, 0])\n",
    "ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)\n",
    "if xlog:\n",
    "    plt.xscale('log')\n",
    "# plt.yscale('log')\n",
    "ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)\n",
    "\n",
    "# plt.xscale('log')\n",
    "x0 = np.percentile(x,75)\n",
    "x1 = np.max(x)\n",
    "# x1 = 10**5\n",
    "y1 = np.max(y)\n",
    "y1 = custom_ceil(y1,50)\n",
    "ax.add_patch(\n",
    "     patches.Rectangle(\n",
    "        (x0, 0),\n",
    "        (x1-x0),\n",
    "        y1,\n",
    "        edgecolor = None,\n",
    "        facecolor = [.2]*3 if darkmode else [.8]*3,\n",
    "        fill=True,\n",
    "        zorder=1\n",
    "     ))\n",
    "# ax.vlines(x0,0,y1,colors=[.5]*4,linestyles='dashed',linewidth=1)\n",
    "# use the previously defined function to do scatter plot\n",
    "sel = x>=5 # remove the 10ish 1-5px mistakes in GT\n",
    "scatter_hist(x[sel], y[sel], ax, ax_histx, ax_histy, zorder=2, hist=0)\n",
    "# plt.xscale('log')\n",
    "# plt.yscale('log')\n",
    "ax.set_ylabel(names[j]+' \\n Segmentation errors', fontsize = labelsize)\n",
    "ax.set_xlabel('Cell area (px$^2$)', fontsize = labelsize)\n",
    "ax.set_ylim(0,y1)\n",
    "ax.set_xlim(1e0,custom_ceil(x1,10**4.3))\n",
    "ax.patch.set_alpha(0.0)\n",
    "\n",
    "fig.patch.set_facecolor(background_color)\n",
    "\n",
    "# set x ticks\n",
    "if xlog:\n",
    "    # plt.xscale('log')\n",
    "    x_major = mpl.ticker.LogLocator(base = 10.0, numticks = 10)\n",
    "    ax.xaxis.set_major_locator(x_major)\n",
    "    x_minor = mpl.ticker.LogLocator(base = 10.0, subs = np.arange(1.0, 10.0) * 0.1, numticks = 10)\n",
    "    ax.xaxis.set_minor_locator(x_minor)\n",
    "    ax.xaxis.set_minor_formatter(mpl.ticker.NullFormatter())\n",
    "\n",
    "# set y ticks\n",
    "# y_major = mpl.ticker.LinearLocator(numticks = 1+5)\n",
    "# ax.yaxis.set_major_locator(y_major)\n",
    "# y_minor = mpl.ticker.LogLocator(base = 10.0, subs = np.arange(1.0, 10.0) * 0.1, numticks = 10)\n",
    "# ax.yaxis.set_minor_locator(y_minor)\n",
    "# ax.yaxis.set_minor_formatter(mpl.ticker.NullFormatter())\n",
    "\n",
    "ax.tick_params(axis='both', which='major', labelsize=labelsize)\n",
    "\n",
    "# x_major = matplotlib.ticker.LogLocator(base = 10.0, numticks = 10)\n",
    "# ax_histy.yaxis.set_major_locator(x_major)\n",
    "# x_minor = matplotlib.ticker.LogLocator(base = 10.0, subs = np.arange(1.0, 10.0) * 0.1, numticks = 10)\n",
    "# ax_histy.yaxis.set_minor_locator(x_minor)\n",
    "# ax_histy.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())\n",
    "\n",
    "# turn off hists\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax_histy.axis('off')\n",
    "ax_histx.axis('off')\n",
    "\n",
    "# ax_histy.spines['right'].set_visible(False)\n",
    "# ax_histy.spines['top'].set_visible(False)\n",
    "# ax_histy.spines['bottom'].set_visible(False)\n",
    "# ax_histy.spines['left'].set_visible(False)\n",
    "# ax_histy.get_xaxis().set_ticks([])\n",
    "# ax_histy.get_yaxis().set_ticks([])\n",
    "# ax_histy.get_yaxis().set_visible(False)\n",
    "# plt.xlim([xmin,xmax])\n",
    "\n",
    "plt.show()\n",
    "for ext in ['.pdf','.eps']:\n",
    "    fig.savefig(os.path.join(save0,'errors_vs_area_hist'+suffix+'_'+names[j]+ext),bbox_inches='tight') #<<<\n",
    "\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/errors_vs_area_hist'+suffix+'.eps',bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/errors_vs_area_hist'+suffix+'_'+names[j]+'.pdf',bbox_inches='tight') #<<<\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/errors_vs_area_hist.png',bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Extended Data Figure 4b\n",
    "\n",
    "This portion of the script depends on the bact_phase dataset being loaded. It could be adapted fairly easily for the other datasets with alternative choices for image names and cell indexes.\n",
    "\n",
    "## Select examples of errored cells\n",
    "From inspection of the test image set, I defined the `img_list` with candidates for making a good figure. This is entirely qualitative, but I wanted to make sure that I included both branched and unbranched morphologies as well as over-segmentation by discrete clustering and over-segmentation in extended clustering. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from scipy.ndimage import binary_erosion, binary_dilation\n",
    "from cellpose_omni import utils, transforms\n",
    "from omnipose.utils import sinebow\n",
    "import ncolor\n",
    "import cv2\n",
    "from cellpose_omni.io import imsave\n",
    "from PIL import Image, ImageFont, ImageDraw, ImageOps\n",
    "from matplotlib.colors import ListedColormap\n",
    "import os, datetime, gc, warnings, glob\n",
    "\n",
    "\n",
    "\n",
    "plt.style.use('dark_background')\n",
    "cmap = mpl.cm.get_cmap('viridis')\n",
    "cm2 = ListedColormap([color for color in sinebow(5).values()])\n",
    "cm3 = ListedColormap([color for color in sinebow(5,bg_color=[1,1,1,1]).values()])\n",
    "\n",
    "# basedir = '/home/kcutler/DataDrive/omnipose_paper/Comparison Examples/newcompare4/'\n",
    "# io.check_dir(basedir)\n",
    "\n",
    "area_cutoff = 20\n",
    "error_cutoff=2\n",
    "pad = 10\n",
    "N = len(masks)\n",
    "nimg = len(masks_gt)\n",
    "bkct = 500\n",
    "bg = 0.5\n",
    "ext = '.png'\n",
    "\n",
    "cnt = 0\n",
    "txtpad = 10\n",
    "yoffset = [0,0]\n",
    "\n",
    "img_list = ['ftsN_ensemble_0','caulo_15','streptomyces_XY15_1','Az_branch_ec_0','Hpylori2_2',\n",
    "            'PSVB_ensemble_c_8','PSVB_ensemble_c_5','PSVB_ensemble_c_2','PSVB_ensemble_c_0','PSVB_ensemble_c_12','PSVB_ensemble_c_11','PSVB_ensemble_c_10',\n",
    "            'vibrio_ensemble_2_19','wiggins_ensemble_1','wiggins_ensemble_11']\n",
    "\n",
    "# cells were chosen by using the Matlotlib Widget and hovering over cells \n",
    "# that I wanted to extract to get the ground truth index \n",
    "# e.g. plt.imshow(masks_gt_clean[img_index[6]])\n",
    "cell_list = [10,8,None,33,44,\n",
    "             160,193,31,58,12,49,10,#[180,194,201]\n",
    "             None,399,119]# [370,399]\n",
    "\n",
    "namelist = [os.path.splitext(os.path.split(file)[-1])[0] for file in img_names]\n",
    "img_index = [namelist.index(im) for im in img_list]\n",
    "[print(l,k) for l,k in zip(cell_list,img_index)]\n",
    "\n",
    "n_exmpl = len(img_list)\n",
    "label_list = []*n_exmpl\n",
    "coords = []*n_exmpl\n",
    "\n",
    "for j in [0]: #base selection on cellpose examples\n",
    "    for l,k in zip(cell_list,img_index): # replace loops over k and l \n",
    "        \n",
    "        file = img_names[k]\n",
    "        basename = os.path.splitext(os.path.split(file)[-1])[0]\n",
    "\n",
    "        mgt = masks_gt_clean[k]\n",
    "        p = transforms.normalize99(imgs[k],omni=True)\n",
    "        img0 = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mgt==0)])))\n",
    "\n",
    "        if l is not None:\n",
    "            if isinstance(l, list):\n",
    "                hits = np.any(np.stack(([mgt==li+1 for li in l])),axis=0)\n",
    "            else:\n",
    "                hits = mgt==l+1\n",
    "            \n",
    "            microcolonies = mgt>0\n",
    "            labels = skimage.measure.label(microcolonies)\n",
    "            binmask = hits.copy()\n",
    "            for cell_ID in np.unique(labels[labels>0]):\n",
    "                mask = labels==cell_ID \n",
    "                area = np.count_nonzero(mask)\n",
    "                overlap = np.count_nonzero(np.logical_and(mask, binary_dilation(hits, iterations=1)))\n",
    "                if overlap > 0: #only premove cells that are 50% or more edge px\n",
    "                    binmask[mask] = 1\n",
    "\n",
    "            y,x = np.nonzero(binmask)\n",
    "            max_y,max_x = np.array(mgt.shape)-1\n",
    "\n",
    "            y0 = max(0,min(y)-pad)\n",
    "            y1 = min(max_y,max(y)+pad)\n",
    "            x0 = max(0,min(x)-pad)\n",
    "            x1 = min(max_x,max(x)+pad)\n",
    "\n",
    "            # p = img0[y0:y1,x0:x1]\n",
    "            # mask_gt = mgt[y0:y1,x0:x1]\n",
    "            xy = [y0,y1,x0,x1]\n",
    "        else:\n",
    "            ly,lx = mgt.shape\n",
    "            xy = [0,ly,0,lx]\n",
    "        \n",
    "        coords.append(xy)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the pixel trajectories "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# re-run both models on full images to reproduce the exact reults from main code\n",
    "# (Crop later; cropped images can have slight significant differences in results for very small crops). \n",
    "chans = [0,0] \n",
    "names = ['Cellpose','Omnipose']\n",
    "suffixes = ['_CP','_OM']\n",
    "# imglist = [imgs[k] for k in img_index]\n",
    "imglist = []\n",
    "\n",
    "# as it turns out, the sample I ultimately selected needs to be gamma-normalized\n",
    "# to reproduce the exact flow field in the paper. Omnipose is not sensitive to this since I\n",
    "# have a gamma augmentation, but Cellpose is, and the flow changes a lot based on that \n",
    "bg = 0.5\n",
    "for k in img_index:\n",
    "    mgt = masks_gt[k].copy()\n",
    "    p = transforms.normalize99(imgs[k],omni=True)\n",
    "    img0 = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mgt==0)])))\n",
    "    imglist.append(img0)\n",
    "# imglist = [imgs[k] for k in [img_index[4]]]\n",
    "\n",
    "J = len(names)\n",
    "nimg = len(imgs)\n",
    "masks,flows,styles,d = [[]]*J,[[]]*J,[[]]*J,[[]]*J\n",
    "for j in range(J):\n",
    "# for j in [1]:\n",
    "    masks[j], flows[j], styles[j], = model[j].eval(imglist,channels=[0,0], mask_threshold=-1, diameter=0,cluster=j,\n",
    "                                                   flow_threshold=0, omni=j, calc_trace=True, min_size=0, tile=False,\n",
    "                                                   transparency=transparency, verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# show the results\n",
    "%matplotlib inline\n",
    "from cellpose_omni import plot\n",
    "j = 0\n",
    "nimg = len(imgs)\n",
    "for idx,im in enumerate(imglist):\n",
    "    maski = masks[j][idx]\n",
    "    flowi = flows[j][idx][0]\n",
    "    fig = plt.figure(figsize=(12,5))\n",
    "    plot.show_segmentation(fig, im, maski, flowi, channels=chans, omni=True, bg_color=0)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure 2b\n",
    "\n",
    "This generates our figure demonstrating how Cellpose over-segments cells. `do_flows` toggles whether or not to do the time-consuming flow line vector graphic showing how boundary pixels coalesce. `isolated` toggles whether or not the panels generate that information only for selected cells or for all cells in the FoV. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import edt\n",
    "from matplotlib import rc\n",
    "from omnipose import utils\n",
    "import ncolor\n",
    "import cmasher as cmr\n",
    "from skimage import filters\n",
    "\n",
    "subdir = 'error_examples'\n",
    "# io.check_dir(subdir)\n",
    "\n",
    "darkmode=0\n",
    "do_flows = 1\n",
    "isolated = 0\n",
    "recompute = 0\n",
    "\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42\n",
    "mpl.rcParams['font.family'] = 'sans-serif'\n",
    "rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})\n",
    "import matplotlib_inline\n",
    "# matplotlib_inline.backend_inline.set_matplotlib_formats('retina', 'png')\n",
    "mpl.rcParams['figure.dpi'] = 72\n",
    "%matplotlib inline\n",
    "A = 1\n",
    "px = A/plt.rcParams['figure.dpi']  # pixel in inches\n",
    "# cmap = mpl.cm.get_cmap('plasma')\n",
    "cmap = cmr.get_sub_cmap('plasma', 0, 0.95)\n",
    "cmap2 = cmr.get_sub_cmap('gray', 0, 1)\n",
    "\n",
    "\n",
    "\n",
    "for j in range(J): # select from models\n",
    "# for j in [-1]:\n",
    "    suffix = suffixes[j]\n",
    "    omni = 'OM' in suffix\n",
    "    print(suffix,omni)\n",
    "    # for k in range(n_exmpl): # select from images \n",
    "    # for k in [4]:\n",
    "    for k in [3,4]:\n",
    "    \n",
    "    # for k in [6]: # limit to this particular cell \n",
    "    \n",
    "        xy = coords[k].copy() # for cropping\n",
    "        ly,lx = xy[1]-xy[0],xy[3]-xy[2]\n",
    "        \n",
    "        # outline thickness\n",
    "        if px*lx<2:\n",
    "            mode = 'inner'\n",
    "        else:\n",
    "            mode = 'thick'\n",
    "        \n",
    "        mgt = masks_gt[img_index[k]].copy()\n",
    "        # p = transforms.normalize99(imglist[k],omni=True)\n",
    "        # img0 = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mgt==0)])))\n",
    "        img0 = imglist[k]\n",
    "        \n",
    "        mpred = masks[j][k].copy()\n",
    "        dP_pred = flows[j][k][1].copy()\n",
    "        cellprob = flows[j][k][2].copy()\n",
    "        bd = flows[j][k][-2].copy() if omni else np.zeros_like(cellprob)\n",
    "        bd = 1/(1+np.exp(-bd))\n",
    "        plt.imshow(utils.rescale(bd))\n",
    "        plt.show()\n",
    "        tr = flows[j][k][-1][0].copy()\n",
    "\n",
    "        # make veritcal\n",
    "        if ly<lx:\n",
    "            print('transposing')\n",
    "            mgt = np.transpose(mgt)\n",
    "            img0 =  np.transpose(img0)\n",
    "            dP_pred = np.transpose(dP_pred,(0,2,1))\n",
    "            #traspose\n",
    "            v1 =  dP_pred[1].copy()\n",
    "            v2 =  dP_pred[0].copy()\n",
    "            theta = 90\n",
    "            dP_pred[0] = (-v1 * np.sin(-theta) + v2*np.cos(-theta))\n",
    "            dP_pred[1] = -(v1 * np.cos(-theta) + v2*np.sin(-theta))\n",
    "            tr = np.stack((tr[1],tr[0]))\n",
    "            cellprob = np.transpose(cellprob)\n",
    "            bd = np.transpose(bd)\n",
    "            mpred = np.transpose(mpred)\n",
    "            xy = [xy[2],xy[3],xy[0],xy[1]]\n",
    "        \n",
    "        #crop\n",
    "        mpred_full = mpred.copy()\n",
    "        mpred = mpred[xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        mgt = mgt[xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        cellprob = cellprob[xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        bd = bd[xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        dP_pred = dP_pred[:,xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        img0 = img0[xy[0]:xy[1],xy[2]:xy[3]]\n",
    "        match = np.argwhere([np.logical_and.reduce((xy[3]>=tr[1,i,0],xy[2]<=tr[1,i,0],xy[0]<=tr[0,i,0],xy[1]>=tr[0,i,0])) for i in range(tr.shape[1])]).flatten()\n",
    "        tr = np.stack(([tr[0,i,:]-xy[0] for i in match],[tr[1,i,:]-xy[2] for i in match]))\n",
    "        # for some reason, the ones I selected came from running the algorithm on\n",
    "        # the cropped field of view. This actually changes the output. \n",
    "        \n",
    "        if recompute:\n",
    "            mpred,f,s = model[j].eval(img0,channels=[0,0], mask_threshold=-1, diameter=0,cluster=j,\n",
    "                                                   flow_threshold=0, omni=j, calc_trace=True, min_size=0, tile=False,\n",
    "                                                   transparency=transparency, verbose=verbose)\n",
    "            \n",
    "\n",
    "            dP_pred = f[1].copy()\n",
    "            cellprob = f[2].copy()\n",
    "            bd = f[-2].copy() if omni else np.zeros_like(cellprob)\n",
    "            bd = 1/(1+np.exp(-bd))\n",
    "            tr = f[-1][0].copy()\n",
    "\n",
    "        \n",
    "        \n",
    "        flow_pred = plot.dx_to_circ(dP_pred,transparency=transparency)\n",
    "        l = cell_list[k] # cell label list\n",
    "        \n",
    "        bini = mgt>0\n",
    "        bin0 = np.stack((bini,bini,bini),axis=2)\n",
    "        \n",
    "        savedir = os.path.join(save0, subdir, img_list[k] + '_cell_number_'+ str(l))\n",
    "        io.check_dir(savedir)\n",
    "        path = os.path.join(savedir,'perim_flows' + suffix)\n",
    "        print(path)\n",
    "        \n",
    "        skimage.io.imsave(os.path.join(savedir,'phase'+ext),np.uint8(img0*255))\n",
    "        skimage.io.imsave(os.path.join(savedir,'flow'+suffix+ext),np.uint8(flow_pred))\n",
    "        pic = cmap(utils.rescale(cellprob))\n",
    "        pic[:,:,-1] = utils.rescale(cellprob)\n",
    "        skimage.io.imsave(os.path.join(savedir,'cellprob'+suffix+ext),((pic)*255).astype(np.uint8))\n",
    "        \n",
    "        pic = cmap2(utils.normalize99(bd))\n",
    "        pic[:,:,-1] = utils.rescale(bd)\n",
    "        skimage.io.imsave(os.path.join(savedir,'boundary'+suffix+ext),((pic)*255).astype(np.uint8))\n",
    "        \n",
    "        mask_threshold = -1\n",
    "        if omni:\n",
    "            mask0 = filters.apply_hysteresis_threshold(cellprob, mask_threshold-1, mask_threshold)\n",
    "        else:\n",
    "            mask0 = cellprob>mask_threshold\n",
    "        b = np.zeros_like(mask0)\n",
    "        pic = np.stack((mask0,b,b),axis=-1)\n",
    "        skimage.io.imsave(os.path.join(savedir,'cellprob_thresh'+suffix+ext),((pic)*255).astype(np.uint8))\n",
    "\n",
    "        if isolated:\n",
    "            mag = transforms.normalize99(np.sqrt(np.sum(dP_pred**2,axis=0)),omni=True)\n",
    "            f = flow_pred.copy()\n",
    "            flow_gray = 0.2125*f[:,:,0] + 0.7154*f[:,:,1] + 0.0721*f[:,:,2]\n",
    "            m = bini==0\n",
    "            f[m] = np.stack([flow_gray,flow_gray,flow_gray,mag],axis=-1)[m]\n",
    "            skimage.io.imsave(os.path.join(savedir,'flow_gray'+suffix+ext),np.uint8(f*255))\n",
    "        \n",
    "#         ncolor_gt = utils.ncolor.label(mgt)\n",
    "        m = mgt.copy()\n",
    "        u = np.unique(mgt)\n",
    "        U = len(u)\n",
    "        A = 100\n",
    "        v = [0]+list(np.linspace(.25,.55,U-1)*A)\n",
    "        res = dict(zip(u, v))\n",
    "        m = fastremap.remap(m,res,preserve_missing_labels=False, in_place=True)\n",
    "        ncolor_gt = utils.rescale(ncolor.label(m.copy()))\n",
    "        ncolor_gray = np.stack([utils.rescale(ncolor_gt)]*3,axis=-1)\n",
    "        skimage.io.imsave(os.path.join(savedir,'ncolor_gray_masks'+ext),np.uint8(ncolor_gray*255))\n",
    "\n",
    "        ncolor_gray = np.stack([1-ncolor_gt]*3,axis=-1)\n",
    "        skimage.io.imsave(os.path.join(savedir,'ncolor_gray_masks_inv'+ext),np.uint8(ncolor_gray*255))\n",
    "        \n",
    "        ncolor_pred = utils.rescale(ncolor.label(mpred.copy()))\n",
    "        skimage.io.imsave(os.path.join(savedir,'ncolor_pred'+ext),np.uint8(ncolor_pred*255))\n",
    "        \n",
    "\n",
    "        outli = plot.outline_view(img0,mpred,color=cmap(.85)[:3],mode=mode)\n",
    "        skimage.io.imsave(os.path.join(savedir,'outline_view_gold'+suffix+ext),np.uint8(outli))\n",
    "        outli = plot.outline_view(img0,mpred,mode=mode)\n",
    "        skimage.io.imsave(os.path.join(savedir,'outline_view_red'+suffix+ext),np.uint8(outli))\n",
    "       \n",
    "        mgt, remap = fastremap.renumber(mgt)\n",
    "        iou = metrics._intersection_over_union(mgt, mpred)\n",
    "        th = 0\n",
    "        n_min = min(iou.shape[0], iou.shape[1])\n",
    "        costs = -(iou >= th).astype(float) - iou / (2*n_min)\n",
    "        true_ind, pred_ind = linear_sum_assignment(costs)\n",
    "        miou = iou[true_ind, pred_ind]\n",
    "\n",
    "        # print('match',miou,pred_ind,mask_match,pred_inds,true_ind)\n",
    "        with open(os.path.join(savedir,'MMiou_and_area'+suffix+'.txt'), \"w\") as text_file:\n",
    "            text = f\"Mean Matched IoU: {np.mean(miou)}\\nAreas: {[np.sum(mgt==l) for l in np.unique(mgt)]}\\nThis area: {np.sum(bin0)}\"\n",
    "            print(text, file=text_file)\n",
    "                               \n",
    "        if do_flows:\n",
    "            dists = edt.edt(mpred)\n",
    "            bd = dists==1\n",
    "            \n",
    "            Y,X = np.nonzero(bd)\n",
    "            # Y = Y[np.logical_and(Y>=xy[0],Y<=xy[1])]\n",
    "            # X = X[np.logical_and(X>=xy[2],X<=xy[3])]\n",
    "            a = .5\n",
    "            match0 = [np.any(np.logical_and((X-tr[1,i,0])**2<a,(Y-tr[0,i,0])**2<a)) for i in range(tr.shape[1])]\n",
    "            select_inds = np.argwhere(match0).flatten()\n",
    "            \n",
    "            Y,X = np.nonzero(mpred)\n",
    "            # Y = Y[np.logical_or(Y>=xy[0],Y<=xy[1])]\n",
    "            # X = X[np.logical_or(X>=xy[2],X<=xy[3])]\n",
    "            a = .5\n",
    "            match2 = [np.any(np.logical_and((X-tr[1,i,0])**2<a,(Y-tr[0,i,0])**2<a)) for i in range(tr.shape[1])]\n",
    "            select_inds2 = np.argwhere(match2).flatten()\n",
    "\n",
    "            lx = mpred.shape[1]\n",
    "            ly = mpred.shape[0]\n",
    "\n",
    "            fig,ax = plt.subplots(figsize=(ly*px,lx*px))\n",
    "            # ax.imshow(dists)\n",
    "            ax.set_aspect('equal')\n",
    "            ax.axis('off')\n",
    "            ax.set_position([0, 0, 1, 1])\n",
    "            \n",
    "            for i in select_inds:\n",
    "                xs = tr[1,i,:]\n",
    "                ys = tr[0,i,:]\n",
    "                if darkmode:\n",
    "                    c = [1,1,1,.5]\n",
    "                    s2 = '_darkmode'\n",
    "                else:\n",
    "                    c = [0,0,0,.25]\n",
    "                    s2 = '_lightmode'\n",
    "                ax.plot(xs,ys,c=c,solid_capstyle='round',linewidth=1/2,zorder=1)\n",
    "                \n",
    "            for i in select_inds2:\n",
    "                ax.scatter(tr[1,i,-1],tr[0,i,-1],marker='.',s=px*lx,edgecolor=None,facecolor='r',zorder=2)\n",
    "            \n",
    "            ax.set_xlim([0,lx])\n",
    "            ax.set_ylim([ly,0])\n",
    "            ax.patch.set_alpha(0.)\n",
    "            fig.patch.set_facecolor(None)\n",
    "\n",
    "            plt.show()\n",
    "            \n",
    "            fig.savefig(path+s2+'.pdf',bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "            fig.savefig(path+s2+'.png',bbox_inches='tight',transparent=True,pad_inches=0,dpi=1000)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "k=3\n",
    "j = 0\n",
    "# xy = [xy[2],xy[3],xy[0],xy[1]]\n",
    "bg = 0.5\n",
    "mgt = masks_gt[img_index[k]].copy()\n",
    "p = utils.rescale(imglist[k])\n",
    "# p = imglist[k]\n",
    "img0 = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mgt==0)])))\n",
    "im = (img0)\n",
    "# im = p\n",
    "m,f,s = model[j].eval(im,channels=[0,0], mask_threshold=-1, diameter=0,cluster=j,\n",
    "                                                   flow_threshold=0, omni=j, calc_trace=True, min_size=0, tile=False,\n",
    "                                                   transparency=transparency, verbose=verbose)\n",
    "\n",
    "maski = m\n",
    "flowi = f[0]\n",
    "fig = plt.figure(figsize=(12,12))\n",
    "plot.show_segmentation(fig, im, maski, flowi, channels=chans, omni=True, bg_color=0)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Export Error examples\n",
    "From cell errors we can look at tons of examples. My current code just gives every single instance of an error. We see from the above histogram that my Omnipose barely gives any errors above 1 (these are issues with cell division, I think). \n",
    "\n",
    "- Enforce a minimum crop size \n",
    "- Sort by number of errors \n",
    "- Dump images into a PDF with info "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from scipy.ndimage.morphology import binary_erosion \n",
    "from cellpose_omni import utils, transforms, plot\n",
    "import cv2\n",
    "from cellpose_omni.io import imsave\n",
    "from PIL import Image, ImageFont, ImageDraw, ImageOps\n",
    "from matplotlib.colors import ListedColormap\n",
    "import os, datetime, gc, warnings, glob\n",
    "from fpdf import FPDF\n",
    "\n",
    "plt.style.use('dark_background')\n",
    "cmap = mpl.cm.get_cmap('viridis')\n",
    "cm2 = ListedColormap([color for color in sinebow(5).values()])\n",
    "\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_paper/Comparison Examples/Per_Cell_GT_Comparison/'\n",
    "io.check_dir(basedir)\n",
    "    \n",
    "cutoff=1\n",
    "pad = 10\n",
    "N = len(masks)\n",
    "nimg = len(masks_gt)\n",
    "bkct = 500\n",
    "bg = 0.5\n",
    "\n",
    "\n",
    "txtpad = 10\n",
    "yoffset = [0,0]\n",
    "buffer = 30\n",
    "do_pdf = 0\n",
    "clean = 1\n",
    "# for k in range(nimg):\n",
    "# for k in [0]:\n",
    "for j in [0]:\n",
    "    if do_pdf:\n",
    "        pdf = FPDF(unit='pt')\n",
    "        pdf.add_page()\n",
    "        pdf.set_font('Arial', 'B', 8)\n",
    "        pdf.set_text_color(r=255, g=255, b=255)\n",
    "        pdf.set_fill_color(0,0,0)\n",
    "        ph = pdf.h\n",
    "        pw = pdf.w\n",
    "        pdf.page_break_trigger = ph\n",
    "        pdf.rect(0,0,pw,ph,'F')\n",
    "        pdf.set_xy(0,0)\n",
    "        ind = 0\n",
    "        cnt = 0\n",
    "    \n",
    "    \n",
    "    for k in range(nimg):\n",
    "#     for k in range(50):\n",
    "        file = img_names[k]\n",
    "        basename = os.path.splitext(os.path.split(file)[-1])[0]\n",
    "\n",
    "        label_list = []\n",
    "        err = cell_errors[j][k]      \n",
    "        err_indexes = np.nonzero(err>=cutoff)[0]\n",
    "        num_errors = err_indexes.size\n",
    "        if num_errors > 0:\n",
    "            \n",
    "            for l in err_indexes:\n",
    "                name = basename + '_' + names[j] + '_cell_number_'+ str(l) + '.png'\n",
    "                savepath = os.path.join(basedir, name)\n",
    "                print(j,k,l,name)\n",
    "                \n",
    "                if not os.path.isfile(savepath) or clean: \n",
    "                    mgt = masks_gt_clean[k]\n",
    "                    p = transforms.normalize99(imgs[k],omni=True)\n",
    "                    img0 = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mgt==0)])))\n",
    "\n",
    "                    y,x = np.nonzero(mgt==l+1)\n",
    "                    max_y,max_x = np.array(mgt.shape)-1\n",
    "\n",
    "                    y0 = max(0,min(y)-pad)\n",
    "                    y1 = min(max_y,max(y)+pad)\n",
    "                    x0 = max(0,min(x)-pad)\n",
    "                    x1 = min(max_x,max(x)+pad)\n",
    "\n",
    "                    p = img0[y0:y1,x0:x1]\n",
    "\n",
    "                    pic = [[]]*2\n",
    "                    img0 = np.stack((p,p,p),axis=2)\n",
    "                    maski = masks[j][k][y0:y1,x0:x1]\n",
    "                    flowi = transforms.normalize99(flows[j][k][y0:y1,x0:x1],omni=True)\n",
    "                    outli = transforms.normalize99(plot.outline_view(img0,maski),omni=True)\n",
    "                    pic[0] = np.hstack((img0,outli,flowi))\n",
    "\n",
    "                    maski = mgt[y0:y1,x0:x1]\n",
    "    #                 print(np.unique(maski),l)\n",
    "                    bini = maski==l+1\n",
    "                    bini = np.stack((bini,bini,bini),axis=2)\n",
    "                    outli = transforms.normalize99(plot.outline_view(img0,maski),omni=True)\n",
    "                    pic[1] = np.hstack((cm2(utils.ncolor.label(maski))[:,:,:3],outli,bini))\n",
    "\n",
    "                    img = np.vstack((pic[0],pic[-1]))\n",
    "                \n",
    "                    io.imsave(savepath,np.uint8(img*255))\n",
    "                \n",
    "                if do_pdf:\n",
    "                    im = Image.open(savepath)\n",
    "                    y = pdf.get_y()\n",
    "\n",
    "                    h = ((pw/2)/im.size[0])*im.size[1]\n",
    "                    if yoffset[ind]+h+txtpad+buffer>=ph:\n",
    "                        pdf.add_page()\n",
    "                        pdf.rect(0,0,pw,ph,'F')\n",
    "                        pdf.set_xy(0,0)\n",
    "                        yoffset = [0,0]\n",
    "                        ind = 0\n",
    "                        pdf.set_xy(0,0)\n",
    "                        cnt = 0\n",
    "                    elif cnt==0 or cnt==1:\n",
    "                        pdf.set_xy(ind*pw/2,yoffset[ind])\n",
    "                    else:\n",
    "                        pdf.set_xy(ind*pw/2,yoffset[ind]+txtpad)\n",
    "\n",
    "                    if h<ph:\n",
    "                        pdf.image(savepath,w=pw/2)\n",
    "                    else:\n",
    "                        pdf.image(savepath,h=ph-2*txtpad)\n",
    "                    pdf.cell(30, 10, txt=name)\n",
    "    #                 pdf.cell(30, 10, txt=str(j)+','+str(k)+','+str(l))\n",
    "    #                 pdf.cell(30, txtpad, str(ind))\n",
    "    #                 print(ind,yoffset)\n",
    "                    yoffset[ind] = pdf.get_y()\n",
    "                    ind = (ind+1)%2\n",
    "                    cnt += 1\n",
    "    if do_pdf:         \n",
    "        pdf.output('/home/kcutler/DataDrive/'+names[j]+'_errors.pdf','F')\n",
    "                "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Figure 5f\n",
    "This section is specifically for extracting out example images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy.ndimage import binary_dilation\n",
    "def get_bbxs(masks,seeds,pad=10):\n",
    "    '''\n",
    "    Get bounding boxes that cover the seed masks.\n",
    "    Takes label matrix masks, seed list seeds. Seeds can be nested list or [None]. \n",
    "    '''\n",
    "    K = len(masks)\n",
    "    bboxes = [[]]*K\n",
    "    for k in range(K): \n",
    "        mgt = masks[k]\n",
    "        microcolonies = mgt>0        \n",
    "        l = seeds[k]\n",
    "        labels = skimage.measure.label(microcolonies)\n",
    "        print(l)\n",
    "        if l is not None:\n",
    "            if isinstance(l, list):\n",
    "                hits = np.any(np.stack(([mgt==li for li in l])),axis=0)\n",
    "            else:\n",
    "                hits = mgt==l\n",
    "            binmask = hits.copy()\n",
    "            for cell_ID in np.unique(labels[labels>0]):\n",
    "                mask = labels==cell_ID \n",
    "                area = np.count_nonzero(mask)\n",
    "                overlap = np.count_nonzero(np.logical_and(mask, binary_dilation(hits, iterations=1)))\n",
    "                if overlap > 0: #only premove cells that are 50% or more edge px\n",
    "                    binmask[mask] = 1\n",
    "\n",
    "            y,x = np.nonzero(binmask)\n",
    "            max_y,max_x = np.array(mgt.shape)-1\n",
    "\n",
    "            y0 = max(0,min(y)-pad)\n",
    "            y1 = min(max_y,max(y)+pad)\n",
    "            x0 = max(0,min(x)-pad)\n",
    "            x1 = min(max_x,max(x)+pad)\n",
    "            # bbox = [y0,y1,x0,x1]\n",
    "            bbox = [y0,x0,y1,x1]\n",
    "        else:\n",
    "            ly,lx = mgt.shape\n",
    "            bbox = [0,0,ly,lx]\n",
    "        bboxes[k] = bbox\n",
    "    return bboxes \n",
    "\n",
    "if dataset=='worm':\n",
    "    suffixes = ['_CP','_MX','_OP','_OP_bact_worm']\n",
    "    subsets = ['wormpose_39','bbc010_44']\n",
    "    cell_seeds = [24,11]\n",
    "\n",
    "elif dataset=='cyto2':\n",
    "    suffixes = ['_CP','_MX','_OP']\n",
    "    subsets = ['001_img']\n",
    "    cell_seeds = [[34,36]]\n",
    "\n",
    "indices = [np.array([i for i,s in enumerate(mask_names) if (name in s)]).astype(int) for name in subsets]\n",
    "indices = [i[0] for i in indices]\n",
    "\n",
    "# I picked the worm labels off of omnipose masks,\n",
    "# cyto2 off of the ground truth\n",
    "if dataset=='worm':\n",
    "    masks_s = [masks[-1][i] for i in indices]\n",
    "else:\n",
    "    masks_s = [masks_gt[i] for i in indices]\n",
    "\n",
    "bboxes = get_bbxs(masks_s,cell_seeds)\n",
    "bboxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import edt\n",
    "from matplotlib import rc\n",
    "from scipy.ndimage import binary_erosion, binary_dilation, zoom\n",
    "import omnipose\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42\n",
    "mpl.rcParams['font.family'] = 'sans-serif'\n",
    "rc('font',**{'family':'sans-serif','sans-serif':['Arial']})\n",
    "import matplotlib_inline\n",
    "# matplotlib_inline.backend_inline.set_matplotlib_formats('retina', 'png')\n",
    "mpl.rcParams['figure.dpi'] = 72\n",
    "%matplotlib inline\n",
    "A = 1\n",
    "px = A/plt.rcParams['figure.dpi']  # pixel in inches\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "outline_col = cmap(0.85)[:3]\n",
    "ext = '.png'\n",
    "bg = 0.5\n",
    "pad = 10\n",
    "\n",
    "def getname(path,suffix=''):\n",
    "    return os.path.splitext(Path(path).name)[0].replace(suffix,'')\n",
    "field_names = [getname(i) for i in img_names]\n",
    "\n",
    "savedir = os.path.join(save0,'examples')\n",
    "io.check_dir(savedir)\n",
    "J = len(suffixes)\n",
    "K = len(masks_s)\n",
    "for j in range(J): # select from models\n",
    "# for j in [-1]:\n",
    "    suffix = suffixes[j]\n",
    "    for i,k in enumerate(indices):\n",
    "        xy = bboxes[i]\n",
    "\n",
    "        ly,lx = xy[1]-xy[0],xy[3]-xy[2]\n",
    "        # outline thickness\n",
    "        val = px*lx\n",
    "        print(val,img_names[k])\n",
    "        if val<3:\n",
    "            mode = 'inner'\n",
    "        else:\n",
    "            mode = 'thick'\n",
    "            \n",
    "            \n",
    "        mgt = masks_gt[k]\n",
    "        p = omnipose.utils.rescale(imglist[k])\n",
    "        slc0 = omnipose.utils.bbox_to_slice(xy,mgt.shape)\n",
    "        slc = (Ellipsis,)+slc0\n",
    "        \n",
    "        img0 = p[slc]\n",
    "        mpred = masks[j][k].copy()[slc]\n",
    "       \n",
    "    \n",
    "        # save yellow/gold outline plots\n",
    "        for c,cn in zip([outline_col,[1,0,0]],['gold','red']):\n",
    "            pic = plot.outline_view(img0,mpred,color=c,mode=mode)\n",
    "            skimage.io.imsave(os.path.join(savedir,field_names[k]+'_outline_'+cn+suffix+ext),pic)\n",
    "        \n",
    "        crop_img = img0\n",
    "        crop_masks = mpred\n",
    "        flow_rgb = flows[j][k][0]\n",
    "        scale = np.array(mgt.shape[0:2])/np.array(flow_rgb.shape[0:2])\n",
    "        crop_flow = zoom(flow_rgb, tuple(scale)+(1,), order=1)[slc0]\n",
    "        \n",
    "        # save RGB flows\n",
    "        savepath = os.path.join(savedir,field_names[k]+'_crop_flows'+suffix+ext)\n",
    "        skimage.io.imsave(savepath,np.uint8(crop_flow))\n",
    "        \n",
    "        # save distance field\n",
    "        savepath = os.path.join(savedir,field_names[k]+'_crop_dist'+suffix+ext)\n",
    "        dist = zoom(flows[j][k][2], tuple(scale), order=1)[slc]\n",
    "        dist = omnipose.utils.rescale(dist)\n",
    "        cmap = mpl.cm.get_cmap('plasma')\n",
    "        pic = cmap(dist)\n",
    "        pic[:,:,-1] = crop_masks>0\n",
    "        skimage.io.imsave(savepath,np.uint8(pic*(2**8-1)))\n",
    "\n",
    "        \n",
    "        \n",
    "        # save the boundary, only for omnipose\n",
    "        bd = flows[j][k][4]\n",
    "        if bd is not None:\n",
    "            savepath = os.path.join(savedir,field_names[k]+'_crop_bd'+suffix+ext)\n",
    "            bd = zoom(bd, tuple(scale), order=1)[slc]\n",
    "            bd = omnipose.utils.rescale(bd)\n",
    "            cmap = mpl.cm.get_cmap('viridis')\n",
    "            pic = cmap(bd)\n",
    "            pic[:,:,-1] = crop_masks>0\n",
    "            skimage.io.imsave(savepath,np.uint8(pic*(2**8-1)))\n",
    "        \n",
    "        \n",
    "        # save a grayscale mask version for adobe illustator vectorization \n",
    "        ncl = ncolor.label(crop_masks)\n",
    "        grey_n = np.stack([1-omnipose.utils.rescale(ncl)]*3,axis=-1)\n",
    "        savepath = os.path.join(savedir,field_names[k]+'_masks_gray'+suffix+ext)\n",
    "        io.imsave(savepath,np.uint8(grey_n*(2**8-1)))\n",
    "        \n",
    "        # save the cropped image, RGB uint8 is not interpolated in illustrator ;) \n",
    "        img0 = transforms.move_min_dim(img0)\n",
    "        if img0.shape[-1] < 3 or img0.ndim < 3:\n",
    "            img0 = plot.image_to_rgb(img0, channels=chans[0] if len(chans)>1 else chans, omni=omni)\n",
    "        savepath = os.path.join(savedir,field_names[k]+'_crop_img'+ext)\n",
    "        io.imsave(savepath,np.uint8(img0))\n",
    "        \n",
    "        if 1:\n",
    "            mgt, remap = fastremap.renumber(mgt[slc])\n",
    "            mpred, remap = fastremap.renumber(mpred)\n",
    "            iou = metrics._intersection_over_union(mgt, mpred)\n",
    "            th = 0\n",
    "            n_min = min(iou.shape[0], iou.shape[1])\n",
    "            costs = -(iou >= th).astype(float) - iou / (2*n_min)\n",
    "            true_ind, pred_ind = linear_sum_assignment(costs)\n",
    "            miou = iou[true_ind, pred_ind]\n",
    "            with open(os.path.join(savedir,'MMiou'+field_names[k]+suffix+'.txt'), \"w\") as text_file:\n",
    "                print(f\"Mean Matched IoU: {np.mean(miou)}\", file=text_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}