scripts/Figure_6_worm_high_res.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Figure 6\n",
    "This noebook extracts a particular example from Luca's high-resolution *C. elegans* dataset. These worms are far larger in terms of pixel diameter than any objects in our training sets, such as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# First, import dependencies.\n",
    "import numpy as np\n",
    "import time, os, sys\n",
    "from cellpose_omni import io, transforms, plot, models, core\n",
    "import omnipose\n",
    "import skimage.io\n",
    "\n",
    "\n",
    "# for plotting \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",
    "\n",
    "# This checks to see if you have set up your GPU properly.\n",
    "# CPU performance is a lot slower, but not a problem if you are only processing a few images.\n",
    "use_GPU = core.use_gpu()\n",
    "print('>>> GPU activated? %d'%use_GPU)\n",
    "\n",
    "# grab specific files\n",
    "# mask = io.imread('/home/kcutler/DataDrive/luca/train/0_masks.png')\n",
    "# img = io.imread('/home/kcutler/DataDrive/luca/train/0.png')\n",
    "\n",
    "basedir = '/home/kcutler/DataDrive/luca/test_manual'\n",
    "mask_filter = '_masks'\n",
    "img_names = io.get_image_files(basedir,mask_filter)\n",
    "mask_names,flow_names = io.get_label_files(img_names, mask_filter)\n",
    "imgs = [skimage.io.imread(f) for f in img_names]\n",
    "masks_gt = [skimage.io.imread(f) for f in mask_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I trained this model with a diameter of 60. Unfortunately, diam_mean is not loaded with the model\n",
    "modeldir = '/home/kcutler/DataDrive/luca/train_good_old/models/cellpose_residual_on_style_on_concatenation_off_omni_train_good_2022_04_21_01_30_25.911493_epoch_3999'\n",
    "diam_mean = 60\n",
    "omni = True\n",
    "model = models.CellposeModel(gpu=use_GPU, pretrained_model=modeldir, net_avg=False, diam_mean=diam_mean, omni=omni)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "diams = [omnipose.core.diameters(mask) for mask in masks_gt]\n",
    "diams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "chans = [0,0]\n",
    "mask_threshold = -1\n",
    "diam_threshold = 12\n",
    "net_avg = 0\n",
    "cluster = 0\n",
    "verbose = 0\n",
    "use_gpu = 1\n",
    "\n",
    "rs = diam_mean/np.array(diams)\n",
    "\n",
    "masks, flows, styles = model.eval(imgs,channels=chans,rescale=rs,mask_threshold=mask_threshold,net_avg=net_avg,\n",
    "                                  transparency=True,flow_threshold=0,omni=omni,resample=True,verbose=verbose,\n",
    "                                  diam_threshold=diam_threshold,cluster=cluster,tile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from cellpose_omni import plot\n",
    "import omnipose\n",
    "nimg = len(imgs)\n",
    "for k in range(nimg):\n",
    "    maski = masks[k]\n",
    "    flowi = flows[k][0]\n",
    "\n",
    "    fig = plt.figure(figsize=(12,5))\n",
    "    plot.show_segmentation(fig, imgs[k], maski, flowi, channels=chans, omni=omni, bg_color=0)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from omnipose import utils\n",
    "from cellpose_omni import metrics\n",
    "from skimage import measure\n",
    "import fastremap\n",
    "\n",
    "threshold=np.linspace(0.5,1,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "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 = utils.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,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "    \n",
    "api,tpi,fpi,fni = metrics.average_precision(masks_gt_clean,masks,threshold=threshold)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(threshold,np.mean(api,axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "axcol = 'k'\n",
    "\n",
    "from omnipose.utils import sinebow\n",
    "names = ['omni']\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,.25,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",
    "j=0\n",
    "ax.plot(threshold,np.mean(api,axis=0),label=names[j],color=colors[j])\n",
    "ax.legend(prop={'size': labelsize/1.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.set_ylabel('Percent of cells \\nwith >=1 error', fontsize = labelsize)\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",
    "plt.ylim([0,1])\n",
    "plt.xlim([.5,1])\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",
    "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",
    "# name = '/home/kcutler/DataDrive/omnipose_paper/Figure 5/error_vs_area_percentiles'+suffix\n",
    "# fig.savefig(name+'.eps',bbox_inches=tight_bbox)\n",
    "# fig.savefig(name+'.pdf',bbox_inches=tight_bbox)\n",
    "# fig.savefig(name+'.eps',pad_inches = 0.5)\n",
    "# fig.savefig(name+'.png',pad_inches = 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ncolor\n",
    "save0 = os.path.join('/home/kcutler/DataDrive/omnipose_paper/Figure 6/luca_worm')\n",
    "io.check_dir(save0)\n",
    "\n",
    "# define outline color\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "outline_col = cmap(0.85)[:3]\n",
    "ext = '.png'\n",
    "mode = 'thick'\n",
    "\n",
    "k = -1 # k denotes the model \n",
    "pad = int(diam/2)\n",
    "n = [0]\n",
    "mgt = mask\n",
    "bin0 = mgt>0\n",
    "\n",
    "inds = np.nonzero(bin0)\n",
    "max_inds = np.array(bin0.shape)-1\n",
    "slc = tuple([slice(max(0,min(inds[k])-pad),min(max_inds[k],max(inds[k])+pad)) for k in range(mgt.ndim)])\n",
    "\n",
    "crop_img = img[(Ellipsis,)*(img.ndim-2)+slc]\n",
    "crop_masks = masks[slc]\n",
    "crop_flow = flows[0][slc]\n",
    "\n",
    "crop_outli = plot.outline_view(crop_img,crop_masks,color=outline_col,mode=mode)\n",
    "\n",
    "plt.imshow(crop_outli,interpolation='none')\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "basedir = os.path.join(save0)\n",
    "io.check_dir(basedir)\n",
    "# save the cropped image, RGB uint8 is not interpolated in illustrator ;) \n",
    "img0 = np.stack([omnipose.utils.normalize99(crop_img),]*3,axis=-1)\n",
    "savepath = os.path.join(basedir,'crop_img'+ext)\n",
    "io.imsave(savepath,np.uint8(img0*(2**8-1)))\n",
    "\n",
    "# save the outlines\n",
    "savepath = os.path.join(basedir,'crop_outlines'+ext)\n",
    "io.imsave(savepath,crop_outli)\n",
    "\n",
    "# save the masks\n",
    "savepath = os.path.join(basedir,'crop_masks'+ext)\n",
    "io.imsave(savepath,np.uint8(crop_masks))\n",
    "\n",
    "# save the flows\n",
    "savepath = os.path.join(basedir,'crop_flows'+ext)\n",
    "skimage.io.imsave(savepath,np.uint8(crop_flow))\n",
    "\n",
    "# save the distance\n",
    "savepath = os.path.join(basedir,'crop_dist'+ext)\n",
    "dist = omnipose.utils.rescale(flows[2][slc])\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",
    "# save the boundary\n",
    "savepath = os.path.join(basedir,'crop_bd'+ext)\n",
    "dist = omnipose.utils.rescale(flows[4][slc])\n",
    "cmap = mpl.cm.get_cmap('viridis')\n",
    "pic = cmap(dist)\n",
    "pic[:,:,-1] = crop_masks>0\n",
    "skimage.io.imsave(savepath,np.uint8(pic*(2**8-1)))\n",
    "\n",
    "#save a grayscale 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(basedir,'masks_gray'+ext)\n",
    "io.imsave(savepath,np.uint8(grey_n*(2**8-1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "crop_img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "img0.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}