scripts/Figure_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "## Figure 1\n",
    "Comparison of existing cell segmentation algorithms on our bacterial dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up document and import dependencies. \n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "from cellpose_omni import plot, models, core, io, dynamics\n",
    "import omnipose, ncolor\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\n",
    "from tifffile import imread\n",
    "use_GPU = core.use_gpu()\n",
    "print('>>> GPU activated? %d'%use_GPU)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "### Load Files\n",
    "We start with loading the ground truth files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "basedir = '/home/kcutler/DataDrive/omnipose_all/phase/test_sorted' # replace this with wherever you put your data \n",
    "mask_filter = '_masks'\n",
    "img_names = io.get_image_files(basedir,mask_filter,look_one_level_down=True)\n",
    "mask_names = io.get_label_files(img_names, mask_filter)\n",
    "\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 = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in mask_names]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "All algorithms were run in their own scripts (also provided), sometimes requiring a specific conda environment described in their documentation. Here we are loading in the results from their various directories. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "basedir = '/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "stardist_dir = os.path.join(basedir,'stardist')\n",
    "paths = sorted(Path(stardist_dir).rglob(\"*masks*\"),key=getname)\n",
    "sd_names = [str(p) for p in paths]\n",
    "masks_sd = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in sd_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "cellpose_dir = os.path.join(basedir,'cellpose','masks')\n",
    "paths = sorted(Path(cellpose_dir).rglob(\"*masks*\"),key=lambda x: getname(x,'_cp_masks'))\n",
    "cp_names = [str(p) for p in paths]\n",
    "masks_cp = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in cp_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "superseg_dir = os.path.join(basedir,'supersegger')\n",
    "paths = sorted(Path(superseg_dir).rglob(\"*masks*\"),key=getname)\n",
    "ss_names = [str(p) for p in paths]\n",
    "masks_ss = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in ss_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrcnn_dir = os.path.join(basedir,'maskrcnn')\n",
    "paths = sorted(Path(mrcnn_dir).rglob(\"*masks*\"),key=getname)\n",
    "mr_names = [str(p) for p in paths]\n",
    "masks_mr = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in mr_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "morpho_dir = os.path.join(basedir,'morphometrics')\n",
    "paths = sorted(Path(morpho_dir).rglob(\"*tif*\"),key=lambda x: getname(x,'_Gparent'))\n",
    "mm_names = [str(p) for p in paths]\n",
    "masks_mm = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in mm_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "#including this in here because we want it for Fig. 5 too\n",
    "omni_dir = os.path.join(basedir,'omnipose','masks')\n",
    "paths = sorted(Path(omni_dir).rglob(\"*masks*\"),key=lambda x: getname(x,'_cp_masks'))\n",
    "om_names = [str(p) for p in paths]\n",
    "masks_om = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in om_names]\n",
    "# masks_om = [skimage.io.imread(f) for f in om_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "misic_train_dir = os.path.join(basedir,'misic_trained')\n",
    "paths = sorted(Path(misic_train_dir).rglob(\"*masks*\"),key=getname)\n",
    "mt_names = [str(p) for p in paths]\n",
    "masks_mt = [ncolor.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in mt_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DeLTA outline overlays were sent to us by Owen for the original 118 images in the test dataset. I've now commented it out but we may go back to it later. \n",
    "\n",
    "# To process these, I need to convert outlines to masks, \n",
    "# which is fine because DeLTA always has at least 1px separation between masks.\n",
    "# Note that the naming convention for the images published on the Cellpose website needed to be numeric, \n",
    "# but thanks to the way I exported those images, they are read in in the correct order here. \n",
    "\n",
    "# delta_dir = os.path.join(basedir,'DeLTA')\n",
    "# paths = sorted(Path(delta_dir).rglob(\"*.png\"),key=getname)\n",
    "# dt_names = [str(p) for p in paths]\n",
    "\n",
    "# from scipy.ndimage import binary_fill_holes \n",
    "# from scipy.ndimage import label\n",
    "# def process_dt(f):\n",
    "#     im = skimage.io.imread(f)/255\n",
    "#     outlines = np.logical_and(im[:,:,0]==1, im[:,:,1]!=1, im[:,:,2]==1)\n",
    "#     binary_masks = binary_fill_holes(outlines)\n",
    "#     masks,_ = label(binary_masks)\n",
    "#     return masks\n",
    "\n",
    "# masks_dt = [ncolor.format_labels(omnipose.utils.clean_boundary(process_dt(f))) for f in dt_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# check that DeLTA matches up with gound-truth\n",
    "# print(len(masks_gt)==len(masks_dt))\n",
    "# n = np.random.randint(len(masks_gt))\n",
    "# plt.imshow(np.hstack((masks_dt[n],masks_gt[n])))\n",
    "# plt.axis('off')\n",
    "nimg = len(masks_om)\n",
    "# np.all([getname(mr_names[k])==getname(cp_names[k],suffix='_cp_masks') for k in range(nimg)])\n",
    "len(masks_om),len(masks_mr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load all masks into list: \n",
    "names = ['Cellpose','StarDist','Mask R-CNN','MiSiC','Morphometrics','SuperSegger','Omnipose']\n",
    "abbrv = ['CP','SD','MR','MS','MM','SS','OM']\n",
    "masks = [masks_cp, masks_sd, masks_mr, masks_mt, masks_mm, masks_ss, masks_om]\n",
    "\n",
    "# have to exlude mask r cnn for now\n",
    "# names = ['Cellpose','StarDist','MiSiC','Morphometrics','SuperSegger','Omnipose']\n",
    "# abbrv = ['CP','SD','MS','MM','SS','OM']\n",
    "# masks = [masks_cp, masks_sd, masks_mt, masks_mm, masks_ss, masks_om]\n",
    "# names = ['Mask R-CNN']\n",
    "# abbrv = ['MR']\n",
    "# masks = [masks_mr]\n",
    "plotc = np.array([[237,28,36],\n",
    "                  [221,130,39],\n",
    "                  [61,177,154],\n",
    "                  [221,206,65],\n",
    "                  [194,105,170],\n",
    "                  [114,125,188],\n",
    "                  [0,0,0]])/255\n",
    "\n",
    "# # masks = [masks_cp,masks]\n",
    "print([len(m) for m in masks],len(names),len(abbrv),len(masks),len(plotc),len(masks_gt))\n",
    "# masks = [masks_cp,masks_om]\n",
    "# masks = [masks_cp,masks]\n",
    "# print(len(masks_gt),len(masks_cp),len(masks_sd),len(masks_ss),len(masks_mh),len(masks_ms),len(masks_om))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "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,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This code ignores image boundary cells. \n",
    "\n",
    "ap = [[]] * len(masks) # average precision matrix per image\n",
    "tp = [[]] * len(masks)\n",
    "fp = [[]] * len(masks)\n",
    "fn = [[]] * len(masks)\n",
    "IoU = [[]] * len(masks) # Intersection over Union \n",
    "OvR = [[]] * len(masks)\n",
    "pred_areas = [[]] * len(masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "nimg = len(masks_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_areas = [[]] * nimg\n",
    "masks_gt_clean = [None]*nimg\n",
    "masks_pred_clean = [[None]*nimg]*len(masks) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# remapping = [[]] * nimg\n",
    "for j in range(nimg):\n",
    "    mgt = ncolor.format_labels(omnipose.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,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify that all the predictions match up with the ground truth \n",
    "for j in range(nimg):\n",
    "    for m,n in zip(masks,names):\n",
    "        if masks_gt_clean[j].shape!=m[j].shape:\n",
    "            print('mismatch:',j,n,masks_gt_clean[j].shape,m[j].shape,masks_gt[j].shape)\n",
    "            print(n)\n",
    "            print(getname(mask_names[j]))\n",
    "            print(getname(mr_names[j]),'\\n')\n",
    "            plt.imshow(np.hstack((masks_gt_clean[j],m[j])))\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "savedir = '/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison/algorithm_performance'\n",
    "io.check_dir(savedir)\n",
    "clean = 0\n",
    "if clean:\n",
    "    # go over each model\n",
    "    for j,masks_pred in enumerate(masks):\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",
    "        print(names[j])\n",
    "        for k in range(nimg):\n",
    "            # get the IoU matrix; axis 0 corresponds to GT, axis 1 to pred \n",
    "\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",
    "\n",
    "            ovp = metrics._label_overlap(masks_gt_clean[k], masks_pred[k])[1:,1:] #throw out columns corresponding to zero  \n",
    "\n",
    "            OvR[j] = OvR[j]+[ovp / areas[np.newaxis,:]] # Overlap Ratio    \n",
    "\n",
    "            IoU[j] = IoU[j]+[iou]\n",
    "\n",
    "        nm = abbrv[j]+'_'+str(j)\n",
    "        np.savez(savedir+'OvR_'+nm,OvR[j])\n",
    "        np.savez(savedir+'IoU_'+nm,IoU[j])\n",
    "        np.savez(savedir+'cell_areas_'+nm,cell_areas[j])\n",
    "        np.savez(savedir+'ap_'+nm,ap[j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save output\n",
    "if clean:\n",
    "    np.savez(os.path.join(savedir,'OvR'),OvR)\n",
    "    np.savez(os.path.join(savedir,'IoU'),IoU)\n",
    "    np.savez(os.path.join(savedir,'cell_areas'),cell_areas)\n",
    "    np.savez(os.path.join(savedir,'ap'),ap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load back in \n",
    "OvR = np.load(os.path.join(savedir,'OvR'+'.npz'),allow_pickle=True)['arr_0']\n",
    "IoU = np.load(os.path.join(savedir,'IoU'+'.npz'),allow_pickle=True)['arr_0']\n",
    "cell_areas = np.load(os.path.join(savedir,'cell_areas'+'.npz'),allow_pickle=True)['arr_0']\n",
    "# remapping = np.load(savedir+'remapping'+'.npz',allow_pickle=True)['arr_0']\n",
    "ap = np.load(os.path.join(savedir,'ap'+'.npz'),allow_pickle=True)['arr_0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "#split dataset into different categories \n",
    "\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",
    "exclude = 'xxxx' # could explude some problematic images for all \n",
    "\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((3,1))\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",
    "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",
    "counts,np.sum(counts),len(mask_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "cc = 0\n",
    "for i in indices[-1]:\n",
    "    ct = len(np.unique(masks_gt_clean[i]))-1\n",
    "    cc += ct\n",
    "    print(getname(mask_names[i]), ct)\n",
    "cc - (1+21+11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "\n",
    "per_cell = 1\n",
    "J = len(IoU)\n",
    "K = len(cats)\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": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_size(w,h, ax=None):\n",
    "    \"\"\" w, h: width, height in inches \"\"\"\n",
    "    if not ax: ax=plt.gca()\n",
    "    l = ax.figure.subplotpars.left\n",
    "    r = ax.figure.subplotpars.right\n",
    "    t = ax.figure.subplotpars.top\n",
    "    b = ax.figure.subplotpars.bottom\n",
    "    figw = float(w)/(r-l)\n",
    "    figh = float(h)/(t-b)\n",
    "    ax.figure.set_size_inches(figw, figh)\n",
    "\n",
    "from omnipose.utils import sinebow\n",
    "J = len(IoU)\n",
    "linestyle='-'\n",
    "# master_color_scheme = ['r','b','g']\n",
    "\n",
    "%matplotlib inline\n",
    "darkmode = 1\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(J)]\n",
    "    background_color = 'k'\n",
    "    suffix = '_dark_mode'\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",
    "x = threshold\n",
    "golden = (1 + 5 ** 0.5) / 2\n",
    "sz = 8\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\n",
    "per_image = 1\n",
    "density = 0\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",
    "    ax = axs[j]\n",
    "    if fill:\n",
    "        err = np.std(ap[j][0],axis=0).T\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[k],linestyle=linestyle,alpha=alpha[k],lw=.75)\n",
    "            ax.plot(x,mean,label=names[j],color='k',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))\n",
    "\n",
    "fig.savefig('/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison/JI_separate.eps',bbox_inches=tight_bbox)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/Fig_1_comparison',bbox_inches=tight_bbox)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/AP_vs_IoU_row.eps',bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/AP_vs_IoU_row.png',bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotc[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig = plt.figure(figsize=(sz, sz/golden)) \n",
    "sz = 1.5\n",
    "fig = plt.figure(figsize=(sz, sz)) \n",
    "ax = plt.axes()\n",
    "\n",
    "\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",
    "#define colors with masks and names\n",
    "colors = sinebow(J*2+2)\n",
    "colors = [colors[2*j+1] for j in range(J)]\n",
    "\n",
    "fill = 0\n",
    "per_image = 1\n",
    "density = 0\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",
    "# colors = ['k',[.5,.5,.5],[.8,.8,.8],'y','r']\n",
    "\n",
    "x0 = .5\n",
    "for j in range(J):\n",
    "    names[j]\n",
    "    c = colors[j] if darkmode else plotc[j]\n",
    "    JI = np.mean(ap[j][0],axis=0).T\n",
    "    ax.plot(x,JI,label=names[j],color=c)\n",
    "    \n",
    "    # mJI = np.sum(np.array([ji*np.sum(a) for a,ji in zip(cell_areas,ap[j][0])]),axis=0)/np.sum(np.concatenate(cell_areas))\n",
    "    # ax.plot(x,mJI.T,'--',label=names[j],color=c,linewidth=.5)\n",
    "    \n",
    "    \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",
    "    \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",
    "    print(names[j],JI[0],x[0],JI[49],x[49],JI[80],x[80])\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((x0,1))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.patch.set_alpha(0.0)\n",
    "\n",
    "# ax.vlines(.8,0,1,colors=[.5]*4,linestyles='dashed',linewidth=1)\n",
    "\n",
    "fig.patch.set_facecolor(background_color)\n",
    "plt.ylabel('Jaccard Index', fontsize = labelsize)\n",
    "plt.xlabel('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",
    "ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.0, 1))\n",
    "# ax.set_yscale('log')\n",
    "# ax.set_ylim(1e-3,1)\n",
    "ax.set_ylim(0,1)\n",
    "# ax.set_xticks([x0,.9,1])\n",
    "# plt.subplots_adjust(bottom=.3,left=0,right=1)\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))\n",
    "\n",
    "fig.savefig('/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison/JI_mean'+suffix+'.eps',bbox_inches=tight_bbox)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/Fig_1_comparison',bbox_inches=tight_bbox)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/AP_vs_IoU_row.eps',bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/AP_vs_IoU_row.png',bbox_inches=\"tight\",pad_inches = 0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "mpl.rcParams[\"text.usetex\"] = True\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42\n",
    "mpl.rcParams['svg.fonttype'] = 'none'\n",
    "mpl.rcParams.update({\n",
    "    \"text.usetex\": False,########################\n",
    "    \"font.family\": \"sans-serif\",\n",
    "    \"font.sans-serif\": [\"Helvetica\"]})\n",
    "# matplotlib_inline.backend_inline.set_matplotlib_formats('retina', 'png')\n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "mpl.rcParams[\"ps.usedistiller\"] = 'xpdf'\n",
    "mpl.rc('text.latex', preamble=r'\\usepackage{cmbright}')\n",
    "\n",
    "# plt.hist(IoU[1][indices[1]].transpose())\n",
    "# plt.show()\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.markers import MarkerStyle\n",
    "from matplotlib.lines import Line2D\n",
    "import seaborn as sns\n",
    "np.random.seed(123)\n",
    "w0 = 0.7\n",
    "split=1\n",
    "box = False\n",
    "scatter = 0\n",
    "violin = 1\n",
    "\n",
    "J = len(IoU)\n",
    "\n",
    "# sz = 6.5\n",
    "sz = 6.5\n",
    "fig, axs = plt.subplots(3,3,figsize=(sz,sz))\n",
    "plt.subplots_adjust(wspace=.4, hspace = .4)\n",
    "axs = axs.flatten()\n",
    "# plot AP vs IoU\n",
    "c = sinebow(J+1)\n",
    "APcolors = [c[j+1] for j in range(J)]\n",
    "\n",
    "# cmap = mpl.cm.get_cmap('Accent')\n",
    "# APcolors = cmap(np.linspace(0,1,len(names)))\n",
    "\n",
    "# APcolors = np.array([[1,0,0],[1,.5,.5],[1,.75,.75],[0,0,0],[.75,.75,.75]])\n",
    "# APcolors = np.array([[.25,1,.25],[.75,1,.25],[.85,.85,.25],[0,0,0],[.75,.75,.75]])\n",
    "\n",
    "# cmap1 = np.fliplr(sns.color_palette(\"YlGn_r\", 6))[:3]\n",
    "# APcolors = list(cmap1[:])\n",
    "# APcolors.append('gray')\n",
    "# APcolors.append('b')\n",
    "\n",
    "ax = axs[0]\n",
    "for j in range(J):\n",
    "    names[j]\n",
    "\n",
    "    x = threshold\n",
    "    ax.plot(x,np.mean(ap[j][0],axis=0).T,label=abbrv[j],color=plotc[j])#APcolors[j]\n",
    "\n",
    "ax.vlines(.8,0,1,colors=[.5]*4,linestyles='dashed',linewidth=1)\n",
    "legend = ax.legend(prop={'size': 5}, loc='lower left', ncol=int(J/2),frameon=False, bbox_to_anchor=(0, -.325),borderpad=0,\n",
    "                   labelspacing=.25,borderaxespad=0)\n",
    "# ax.legend(prop={'size': 5}, frameon=True)\n",
    "# for t in legend.get_texts():\n",
    "#     t.set_ha('right') # ha is alias for horizontalalignment\n",
    "#     t.set_position((.1,0))\n",
    "\n",
    "\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,1))\n",
    "ax.set_xticks(np.linspace(.5,1,6))\n",
    "\n",
    "ax.set_yticks(np.linspace(0,1,6))\n",
    "ax.set_yticklabels([0,.2,.4,.6,.8,1])\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.patch.set_alpha(0.0)\n",
    "#     ax.set(adjustable='box', aspect=1/2)\n",
    "ax.set_box_aspect(1)\n",
    "\n",
    "lw = .8\n",
    "\n",
    "# violin plots\n",
    "c = omnipose.utils.sinebow(K)\n",
    "colors = np.array(list(c.values()))[1:]\n",
    "colors = colors.copy()\n",
    "colors[:,-1] = 1\n",
    "colors = ['tab:green','tab:gray','tab:red']\n",
    "\n",
    "cmap = sns.diverging_palette(0, 100, n=3)\n",
    "cmap = sns.color_palette(\"husl\", 8)\n",
    "\n",
    "colors = cmap[:]\n",
    "\n",
    "colors = ['#79C261','#5EACCD','#7B66AD']\n",
    "\n",
    "text = []\n",
    "for j in range(J):\n",
    "    ax = axs[j+1]\n",
    "    \n",
    "    if split:\n",
    "        w = w0 / 3\n",
    "        x = [j-w,j,j+w]\n",
    "    else:\n",
    "        x = [j,j,j]\n",
    "        w = w0\n",
    "    height = [np.mean(yi) for yi in y[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",
    "    \n",
    "    text += [names[j],'Percent above 0.8:\\nBy category: {} \\nOverall: {} \\n'.format(percents,overall_percent)]#,height,len(y[j]))\n",
    "\n",
    "    \n",
    "    parts = ax.violinplot(y[j],positions=x,showextrema=0,widths=w0/3.5)#showmeans=1,quantiles = [[.25,.75]]*len(x),\n",
    "    boxes = ax.boxplot(y[j],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",
    "    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",
    "    # parts['cmeans'].set_color('k')\n",
    "    # parts['cquantiles'].set_color('r')\n",
    "    # parts['cmeans'].set_linewidth(1.5)\n",
    "    # parts['cquantiles'].set_linewidth(1)\n",
    "    # for i, box in enumerate(boxes):\n",
    "    #     box.set_facecolor('w')\n",
    "    # plt.setp(boxes[\"boxes\"], facecolor='w')\n",
    "    \n",
    "    ax.set_ylim((0,1))\n",
    "    ax.set_xlim((j-w0/2,j+w0/2))\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks(np.linspace(0,1,6))\n",
    "    ax.set_yticklabels([0,.2,.4,.6,.8,1])\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",
    "\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.patch.set_alpha(0.0)\n",
    "    ax.set_xlabel(names[j]+' ('+abbrv[j]+')', fontsize = labelsize)\n",
    "    bbx = ax.get_window_extent()\n",
    "    # print(ax.get_aspect(),bbx.height/bbx.width)\n",
    "#     ax.set(adjustable='box', aspect=bbx.width/bbx.height)\n",
    "    ax.set_box_aspect(1)\n",
    "    ax.xaxis.get_major_formatter()._usetex = False\n",
    "    ax.yaxis.get_major_formatter()._usetex = False\n",
    "    xlim = ax.get_xlim()\n",
    "    ax.hlines(.8,xlim[0],xlim[-1],colors=[.5]*4,linestyles='dashed',linewidth=1,zorder=0)\n",
    "\n",
    "#             pc.set_alpha(1)\n",
    "#         ax.vlines(inds, quartile1, quartile3, color='k', linestyle='-', lw=5)\n",
    "\n",
    "# ax.legend(prop={'size': labelsize}, loc='best', frameon=False)\n",
    "# ax.legend(prop={'size': labelsize}, loc='upper left', frameon=False,bbox_to_anchor=(1.0, 1))\n",
    "\n",
    "# ax.set_xticks(range(J))\n",
    "# ax.set_xticklabels(names)\n",
    "\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('Mean Matched IoU', fontsize = labelsize)\n",
    "# # ax.set_xlabel('Algorithm', fontsize = labelsize)\n",
    "\n",
    "ax.patch.set_alpha(0.0)\n",
    "plt.subplots_adjust(wspace=0)\n",
    "\n",
    "\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/fig1_grid.eps',bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "# fig.savefig('/home/kcutler/DataDrive/omnipose_paper/fig1_grid.png',bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison'\n",
    "with open(os.path.join(basedir,'IoU_stats.txt'), \"w\") as text_file:\n",
    "    [print(t,file=f) for t in text for f in [None,text_file]]\n",
    "    \n",
    "fig.savefig(os.path.join(basedir,'violin_grid.pdf'),bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "fig.savefig(os.path.join(basedir,'violin_grid.eps'),bbox_inches=\"tight\",pad_inches = 0.05)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1_names = ['Az_branch_ec_0_','Az_branch_ec_14_','Az_branch_ec_10_','PSVB_ensemble_c_0_', 'PSVB_ensemble_c_1_','streptomyces_XY15_1_']\n",
    "# fig1_names = ['streptomyces_XY15_1_','Az_branch_ec_0_', 'PSVB_ensemble_c_1_']\n",
    "\n",
    "N = len(fig1_names)\n",
    "fig1_indices = np.array([i for i,s in enumerate(mask_names) if any(name in s for name in fig1_names)]).astype(int)\n",
    "fig1_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(names)\n",
    "# for ind in [indices[1][0]]:\n",
    "# for ind in indices[-1]:\n",
    "for ind in fig1_indices:\n",
    "    print(mask_names[ind])\n",
    "    im = imgs[ind]\n",
    "    print('shape',im.shape)\n",
    "    mask_gt = masks_gt[ind]\n",
    "    masks_all = [m[ind] for m in masks]\n",
    "    masks_all.append(mask_gt)\n",
    "    \n",
    "    outli = [plot.outline_view(omnipose.utils.rescale(im),m) for m in masks_all]\n",
    "    \n",
    "    fig = plt.figure(figsize=(12,5))\n",
    "    plt.imshow(np.hstack(outli))\n",
    "#     plt.imshow(im)\n",
    "    plt.axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "#     plt.imshow(mask_gt)\n",
    "#     plt.axis('off')\n",
    "#     plt.tight_layout()\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# crop_dict = {'Az_branch_ec_0_': (372, 163, 170, 160),\n",
    "#              'Az_branch_ec_10_': (268, 236, 180, 120),\n",
    "#              'Az_branch_ec_14_': (418, 297, 190, 120),\n",
    "#              'PSVB_ensemble_c_0_': (330, 350, 160, 230),\n",
    "#              'PSVB_ensemble_c_1_': (0, 0, 190, 390),\n",
    "#              'streptomyces_XY15_1_': (0, 425, 330, 150)}\n",
    "a = 160\n",
    "crop_dict = {'Az_branch_ec_0_': (0, 425, 330, a),\n",
    "             'Az_branch_ec_14_': (372, 163, 175, a),\n",
    "             'Az_branch_ec_10_': (430, 297, a, 120),\n",
    "             'PSVB_ensemble_c_0_': (268, 220, 180, a),\n",
    "             'PSVB_ensemble_c_1_': (330, 350, a, 230),\n",
    "             'streptomyces_XY15_1_': (0, 0, a, 390)}\n",
    "\n",
    "dict_names = list(crop_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(indices[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rgb_pad(im, pad, c=255):\n",
    "    w = 3\n",
    "    im = np.pad(im,((w,w),(w,w),(0,0)),constant_values=0)\n",
    "    return np.pad(im,((pad,pad),(pad,pad),(0,0)),constant_values=c)\n",
    "\n",
    "# cmap = matplotlib.cm.get_cmap('viridis')\n",
    "from matplotlib.colors import ListedColormap\n",
    "from skimage.segmentation import find_boundaries\n",
    "colors = [color for color in sinebow(5).values()]\n",
    "colors[0] = [1,1,1]\n",
    "colormap = ListedColormap(colors)\n",
    "\n",
    "# outline_col = np.array([212,175,55])/255\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "outline_col = cmap(0.85)[:3]\n",
    "\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_all/bact_phase_comparison/seg_examples'\n",
    "io.check_dir(basedir)\n",
    "    \n",
    "%matplotlib inline\n",
    "# produce vector inline graphics\n",
    "# import matplotlib_inline\n",
    "# matplotlib_inline.backend_inline.set_matplotlib_formats('pdf', 'svg')\n",
    "import matplotlib_inline\n",
    "matplotlib_inline.backend_inline.set_matplotlib_formats('retina', 'png')\n",
    "\n",
    "bg = 0.5\n",
    "F = len(dict_names)\n",
    "pad = 5\n",
    "c = 255\n",
    "rows = []\n",
    "# for i in range(F):\n",
    "selection = [1,4,5]\n",
    "# fig, axs = plt.subplots(len(selection), len(names)+1, figsize=(10, 10))#,sharex=True,constrained_layout = True)\n",
    "coords = np.zeros((len(selection), len(names)+1,2))\n",
    "grid_ious = np.zeros((len(selection), len(names)+1))\n",
    "for idx,i in enumerate(selection):\n",
    "    ind = fig1_indices[i]\n",
    "    name = fig1_names[i]\n",
    "    bbx = crop_dict[name]\n",
    "    print(bbx,name)\n",
    "    img0 = omnipose.utils.normalize99(imgs[ind])\n",
    "#     print(img0.shape)\n",
    "    p = img0[bbx[1]:bbx[1]+bbx[3],bbx[0]:bbx[0]+bbx[2]]\n",
    "\n",
    "    \n",
    "    mgt = masks_gt[ind][bbx[1]:bbx[1]+bbx[3],bbx[0]:bbx[0]+bbx[2]]\n",
    "    im = omnipose.utils.normalize_image(p,mgt,bg)\n",
    "    mgt,remap = fastremap.renumber(mgt) \n",
    "    print(np.unique(mgt))\n",
    "    masks_all = []\n",
    "    for m in masks:\n",
    "        m0 = m[ind][bbx[1]:bbx[1]+bbx[3],bbx[0]:bbx[0]+bbx[2]]\n",
    "        m0, remap = fastremap.renumber(m0) \n",
    "        masks_all.append(m0)\n",
    "    ious = [metrics._intersection_over_union(mgt, m)[1:, 1:] for m in masks_all]\n",
    "    for j,iou in enumerate(ious):\n",
    "        \n",
    "        # this metric finds the maximum iou along the GT axis and averages the nonzero ones\n",
    "        # this is biases and allows double-counting \n",
    "        max_iou = np.max(iou,axis=0)\n",
    "        v = np.mean(max_iou[max_iou>0]) \n",
    "        \n",
    "        # instead figure out the best assignment and average that \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",
    "        \n",
    "        v = np.sum(iou[true_ind, pred_ind])/max(iou.shape[0], iou.shape[1]) # instead of average over existing matches, this penalizes matches not found (effectively setting those ious to 0)\n",
    "\n",
    "        print(v,iou.shape,len(true_ind),len(pred_ind))\n",
    "        grid_ious[idx][j] = v\n",
    "        \n",
    "        # \n",
    "        # print('match',miou,pred_ind,mask_match,pred_inds,true_ind)\n",
    "        with open(os.path.join(basedir,'MMiou_cell_'+str(i)+'_'+names[j] +'.txt'), \"w\") as text_file:\n",
    "            print(f\"Mean Matched IoU: {np.mean(v)}\", file=text_file)\n",
    "            \n",
    "    outli_y = [plot.outline_view(im,m,outline_col,mode='thick') for m in masks_all]\n",
    "    outli_r = [plot.outline_view(im,m) for m in masks_all]\n",
    "    \n",
    "    raw_outlines = [find_boundaries(m,mode='inner') for m in masks_all]\n",
    "    color_outlines = [find_boundaries(m,mode='inner')*m for m in masks_all]\n",
    "    \n",
    "    height,width = mgt.shape\n",
    "    flip = height<width and (i != 2)\n",
    "    if flip:\n",
    "        mgt = np.transpose(mgt)\n",
    "        im =  np.transpose(im)\n",
    "    \n",
    "    for k in range(len(outli_y)):\n",
    "        if flip:\n",
    "            outli_y[k] = np.transpose(outli_y[k],(1,0,2))\n",
    "            outli_r[k] = np.transpose(outli_r[k],(1,0,2))\n",
    "            raw_outlines[k] = np.transpose(raw_outlines[k])\n",
    "            color_outlines[k] = np.transpose(color_outlines[k])\n",
    "            masks_all[k] = np.transpose(masks_all[k])\n",
    "        \n",
    "\n",
    "        coords[idx][k][1] = (k+1)*(2*pad+width) - pad\n",
    "        if idx>0:\n",
    "            coords[idx][k][0] = coords[idx-1][k][0] + 2*pad + height\n",
    "        else:\n",
    "            coords[idx][k][0] = pad + height\n",
    "    \n",
    "    K = len(masks_all)\n",
    "    img0 = np.stack([omnipose.utils.rescale(im)]*3,axis=-1)\n",
    "    ncl = ncolor.label(mgt)\n",
    "    grey_n = np.stack([1-omnipose.utils.rescale(ncl)]*3,axis=-1)\n",
    "    color_n = colormap(ncl)[:,:,:3]\n",
    "#     gt_array = np.hstack((rgb_pad(img0, pad),rgb_pad(colormap(omnipose.utils.ncolorlabel(mgt))[:,:,:3],pad)))\n",
    "    gt_array = np.hstack((rgb_pad(img0, pad),rgb_pad(grey_n,pad)))\n",
    "    gt_array = np.pad(gt_array,((0,0),(0,5*pad),(0,0)),constant_values=c) # five spacing between \n",
    "                         \n",
    "    seg_array = np.hstack([rgb_pad(outli_r[k],pad) for k in range(K)])\n",
    "\n",
    "    row = np.hstack([gt_array,seg_array/255])\n",
    "    if idx<len(selection)-1:\n",
    "        row = np.pad(row,((0,5*pad),(0,0),(0,0)),constant_values=c)\n",
    "#     plt.imshow()\n",
    "#     plt.axis('off')\n",
    "#     plt.show()\n",
    "    rows.append(row)\n",
    "    \n",
    "    #saving:\n",
    "   \n",
    "    ext = '.png'\n",
    "    basename = 'cell_'+str(i)\n",
    "    savepath = os.path.join(basedir,basename + '_raw_phase'+ext)\n",
    "    io.imsave(savepath,np.uint8(img0*(2**8-1)))\n",
    "    savepath = os.path.join(basedir,basename + '_masks_color'+ext)\n",
    "    io.imsave(savepath,np.uint8(color_n*(2**8-1)))\n",
    "    savepath = os.path.join(basedir,basename + '_masks_gray'+ext)\n",
    "    io.imsave(savepath,np.uint8(grey_n*(2**8-1)))\n",
    "    for k in range(len(masks_all)): \n",
    "        filename =  basename+'_'+names[k] + '_yellow' + ext\n",
    "        savepath = os.path.join(basedir, filename)\n",
    "        io.imsave(savepath,np.uint8(outli_y[k]))\n",
    "        \n",
    "        filename =  basename+'_'+names[k] + '_red'+ ext\n",
    "        savepath = os.path.join(basedir, filename)\n",
    "        io.imsave(savepath,np.uint8(outli_r[k]))\n",
    "        \n",
    "        filename =  basename+'_color_pred_'+names[k]  + '.eps'\n",
    "        savepath = os.path.join(basedir, filename)\n",
    "        skimage.io.imsave(savepath,np.uint8(colormap(ncolor.label(masks_all[k]))[:,:,:3]*255))\n",
    "        \n",
    "        \n",
    "\n",
    "# outlines = np.array(outlines,dtype=object)\n",
    "grid = np.vstack(rows)\n",
    "grid = np.pad(grid,((pad,pad),(pad,pad),(0,0)),constant_values=c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "ncolor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "scale=1\n",
    "fig = plt.figure(figsize=(scale*grid.shape[1]/72,scale*grid.shape[0]/72))\n",
    "# plt.imshow(cv2.resize((grid).astype(np.uint8)dsize=(scale*grid.shape[1],scale*grid.shape[0]), interpolation=cv2.INTER_NEAREST),interpolation='nearest')\n",
    "# plt.matshow(grid)\n",
    "plt.imshow(grid,interpolation='none')\n",
    "offset = [-1,-1]\n",
    "# for r in range(len(rows))\n",
    "for i in range(coords.shape[0]):\n",
    "    for j in range(coords.shape[1]):\n",
    "# for i in range(coords.shape[0]):\n",
    "#     for j in [0,1]:\n",
    "        y,x = coords[i][j]+[0,30]\n",
    "        x += 2*width+pad\n",
    "        print(x,y)\n",
    "        plt.text(x, y, ('%.3f' % grid_ious[i][j]),color='w',ha = 'right',va='center',fontsize=12) #set rcparams to use font family\n",
    "\n",
    "\n",
    "# fig.axes.get_xaxis().set_visible(False)\n",
    "# fig.axes.get_yaxis().set_visible(False)\n",
    "plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, \n",
    "            hspace = 0, wspace = 0)\n",
    "plt.axis('off')\n",
    "fig.patch.set_facecolor('r')\n",
    "plt.show()\n",
    "\n",
    "#MAYBE use different 4-color scheme/tint for masks based on group \n",
    "# plt.savefig(os.path.join(basedir,'grid.pdf'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = '/home/kcutler/DataDrive/omnipose_paper/qualitative_grid_mean_match_iou'\n",
    "fig.savefig(base+'.eps',bbox_inches=\"tight\",pad_inches = 0.0)\n",
    "fig.savefig(base+'.png',bbox_inches=\"tight\",pad_inches = 0.0)\n",
    "fig.savefig(base+'.pdf',bbox_inches=\"tight\",pad_inches = 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "ws = np.unique([lines.shape[1] for outli in outlines for lines in outli])\n",
    "hs = np.unique([lines.shape[0] for outli in outlines for lines in outli])\n",
    "print(ws,hs)\n",
    "np.lcm.reduce(hs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "basedir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grid_ious)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = [io.imread(f) for f in img_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks_gt_raw = [skimage.io.imread(f) for f in mask_names]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from interactivecrop.interactivecrop import main as crop\n",
    "crop([omnipose.utils.rescale(imgs[i])*255 for i in fig1_indices],image_name_list = fig1_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.transpose(outli[0],(1,0,2)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "outli[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# modeldir_orig = '/home/kcutler/DataDrive/omnipose_train/registered/models/cellpose_residual_on_style_on_concatenation_off_registered_2021_07_22_11_48_00.465748' #500 epochs, bs 8, lr .2\n",
    "# model = models.CellposeModel(gpu=use_GPU, pretrained_model=modeldir_orig)\n",
    "# masks, flows, styles = model.eval(imgs,channels=[0,0],rescale=None,dist_threshold=-1,flow_threshold=0)\n",
    "\n",
    "# basedir = '/home/kcutler/DataDrive/omnipose_paper/Cellpose_StarDist_Etc_Comparison'\n",
    "# if not os.path.isdir(basedir):\n",
    "#     os.mkdir(basedir)\n",
    "# savedir = os.path.join(basedir,'original_cellpose')\n",
    "\n",
    "# if not os.path.isdir(savedir):\n",
    "#     os.mkdir(savedir)\n",
    "# io.save_masks(imgs, masks, flows, img_names,save_flows=True,save_outlines=True,savedir=savedir,png=False,tif=True)\n",
    "# cellpose_dir\n",
    "\n",
    "# stardist_dir = '/home/kcutler/DataDrive/omnipose_paper/Comparison Examples StarDist Bacteria'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in indices[-1]:\n",
    "#     plt.imshow(plot.outline_view(imgs[i],masks_om[i]))\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "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": 5
}