scripts/Figure_3_S5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 3, S5\n",
    "This script pulls out some specific examples of strange cell morphologies to illustrate how Omnipose works. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import time, os, sys\n",
    "from urllib.parse import urlparse\n",
    "import skimage.io\n",
    "\n",
    "import matplotlib as mpl\n",
    "%matplotlib inline\n",
    "\n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap, rgb2hex\n",
    "plt.style.use('dark_background')\n",
    "plt.rcParams['image.cmap'] = 'gray'\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "\n",
    "\n",
    "from cellpose_omni import utils\n",
    "from skimage.io import imread\n",
    "import scipy\n",
    "from skimage.morphology import binary_erosion, binary_dilation, medial_axis, disk, medial_axis,remove_small_holes,remove_small_objects\n",
    "# import skfmm\n",
    "import cv2\n",
    "\n",
    "import fastremap #important for relabeling! Could move just to color function perhaps, but idk of other parts of cellpose require sequential lebel matrices; probably worth it\n",
    "\n",
    "\n",
    "from cellpose_omni import models, io, dynamics, plot, transforms, omnipose\n",
    "from cellpose_omni.omnipose.utils import rescale\n",
    "def laplacian(I):\n",
    "    Iy,Ix = np.gradient(I)\n",
    "    mu = transforms.normalize_field(np.stack((Iy,Ix)))\n",
    "    Iyy,Iyx = np.gradient(mu[0])\n",
    "    Ixy,Ixx = np.gradient(mu[1])\n",
    "    return Ixx+Iyy\n",
    "\n",
    "def divergence(mu):\n",
    "    mask = np.sqrt(np.sum(dP**2,axis=0))\n",
    "    Ly,Lx = mask.shape\n",
    "    Y, X = np.nonzero(mask)\n",
    "    pad = 1\n",
    "    Tx = np.zeros((Ly+2*pad)*(Lx+2*pad), np.float64)\n",
    "    Tx[Y*Lx+X] = np.reshape(dP[1].copy(),Ly*Lx)[Y*Lx+X]\n",
    "    Ty = np.zeros((Ly+2*pad)*(Lx+2*pad), np.float64)\n",
    "    Ty[Y*Lx+X] = np.reshape(dP[0].copy(),Ly*Lx)[Y*Lx+X]\n",
    "\n",
    "    div = np.zeros(Ly*Lx, np.float64)\n",
    "    div[Y*Lx+X]=(Ty[(Y+2)*Lx+X]+8*Ty[(Y+1)*Lx+X]-8*Ty[(Y-1)*Lx+X]-Ty[(Y-2)*Lx+X]+\n",
    "                 Tx[Y*Lx+X+2]+8*Tx[Y*Lx+X+1]-8*Tx[Y*Lx+X-1]-Tx[Y*Lx+X-2])\n",
    "    div.shape = (Ly,Lx)\n",
    "    return div\n",
    "    \n",
    "def derivative(T,masks):\n",
    "    Ly,Lx = T.shape\n",
    "    mu = np.zeros((2,Ly,Lx))\n",
    "    pad = 1\n",
    "    slices = scipy.ndimage.find_objects(masks) \n",
    "    for i,si in enumerate(slices):\n",
    "        if si is not None:\n",
    "            sr,sc = si\n",
    "            mask = np.pad((masks[sr, sc] == i+1),pad)\n",
    "            ly, lx = sr.stop - sr.start + 2*pad, sc.stop - sc.start + 2*pad\n",
    "            y, x = np.nonzero(mask) \n",
    "            ly = np.int32(ly)\n",
    "            lx = np.int32(lx)\n",
    "            y = y.astype(np.int32)  #no need to shift, as array already padded\n",
    "            x = x.astype(np.int32)    \n",
    "\n",
    "            T_ = np.zeros((ly+2*pad)*(lx+2*pad), np.float64)\n",
    "            T_[y*lx+x] = T[sr.start+y-pad, sc.start+x-pad]\n",
    "            dy = (T_[(y+1)*lx + x] - T_[(y-1)*lx + x]) / 2\n",
    "            dx = (T_[y*lx + x+1] - T_[y*lx + x-1]) / 2\n",
    "            mu[:, sr.start+y-pad, sc.start+x-pad] = np.stack((dy,dx))\n",
    "    return mu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = 'cell_4243'\n",
    "# pos = 'cell_4030'\n",
    "# pos = 'hpaz2xy1_10'\n",
    "# pos = 'hpaz2xy1_21'\n",
    "# basedir = '/home/kcutler/DataDrive/omnipose_paper/Figure 3/'+pos\n",
    "\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_paper/Figure 4/'+pos\n",
    "\n",
    "maskdir = basedir + '_mask.png'\n",
    "\n",
    "masks = imread(maskdir).astype(np.uint16)\n",
    "if len(masks.shape)>2:\n",
    "    masks, remap = fastremap.renumber(masks[:,:,1])\n",
    "    \n",
    "bg = 0.5\n",
    "p = omnipose.utils.normalize99(imread(basedir + '_phase.png'))\n",
    "phase = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(masks==0)])))\n",
    "\n",
    "\n",
    "print(\"Loaded array is of type:\", type(masks))\n",
    "print(\"Loaded array has shape:\", masks.shape)\n",
    "print(\"Loaded values are of type:\", masks.dtype)\n",
    "plt.imshow(phase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masks_rgb = np.stack([masks]*3,axis=-1)\n",
    "pic = 1-rescale(masks_rgb)\n",
    "io.imsave(basedir + '_masks_rbg.png',(pic*(2**8-1)).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omni = 1\n",
    "fmm = 0\n",
    "if omni:\n",
    "    if fmm:\n",
    "        suffix = '_fmm'\n",
    "    else:\n",
    "        suffix = ''\n",
    "else:\n",
    "    suffix = '_OG'\n",
    "\n",
    "    \n",
    "if not fmm:\n",
    "    labels,dists,heat,flows = dynamics.masks_to_flows(masks,use_gpu=True,omni=omni)\n",
    "    dP = np.stack([flows[0],flows[1]])\n",
    "else:\n",
    "    heat = edt.edt(masks)\n",
    "    dP = derivative(heat,masks)\n",
    "    dP = transforms.normalize_field(dP)\n",
    "\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "pic = cmap(rescale(heat))\n",
    "pic[:,:,-1] = masks>0\n",
    "plt.imshow(pic)\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "skimage.io.imsave(basedir + '_8bit_heat' + suffix + '.png',((pic)*255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_flow = plot.dx_to_circ(dP,transparency=True)\n",
    "pic = gt_flow\n",
    "plt.imshow(pic)\n",
    "plt.axis('off')\n",
    "skimage.io.imsave(basedir + '_GT_flow'+suffix+'.png',(pic*(2**8-1)).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.imshow(np.sqrt(np.sum(dP**2,axis=0)))\n",
    "div = rescale(divergence(dP))\n",
    "# pic = np.stack([div]*3,axis=-1)\n",
    "pic = cmap(div)\n",
    "# pic[:,:,-1] = rescale(heat**a)\n",
    "pic[:,:,-1] = masks>0\n",
    "plt.imshow(pic)\n",
    "plt.axis('off')\n",
    "skimage.io.imsave(basedir + '_GT_divergence'+suffix+'.png',(pic*(2**8-1)).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omni=1\n",
    "if omni:\n",
    "    if fmm:\n",
    "        suffix = '_fmm'\n",
    "    else:\n",
    "        suffix = ''\n",
    "else:\n",
    "    suffix = '_OG'\n",
    "\n",
    "import edt\n",
    "# inds = np.array(np.nonzero(np.logical_or(masks>0,np.abs(dP[0])>1e-3))).astype(np.int32).T\n",
    "dists = edt.edt(masks)\n",
    "dists[dists==0]=-5\n",
    "inds = np.array(np.nonzero(masks)).astype(np.int32).T\n",
    "print(inds.shape)\n",
    "### update to new code refactoring below\n",
    "masks_pred, p, tr = dynamics.compute_masks(dP,dists,inds=inds,calc_trace=True,use_gpu=0,\n",
    "                                           omni=omni,niter=200,interp=1,mask_threshold=-1,verbose=True,\n",
    "                                          flow_threshold = 0.0,min_size=0)\n",
    "# masks, flows, styles = model[0].eval(phase_list,channels=[0,0],dist_threshold=-1,flow_threshold=0,omni=0, calc_trace=True,min_size=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dists_pred = edt.edt(masks_pred)\n",
    "bd = np.logical_or(dists==1,dists_pred==1)\n",
    "# bd = dists_pred==1\n",
    "plt.imshow(bd)\n",
    "np.any(masks_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import rc\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",
    "\n",
    "A = 1.5\n",
    "px = A/plt.rcParams['figure.dpi']  # pixel in inches\n",
    "\n",
    "print(tr.shape,np.count_nonzero(masks_pred))\n",
    "Y,X = np.nonzero(bd) #boundary only\n",
    "# select_inds = np.flatnonzero(masks)\n",
    "a = .5\n",
    "match = [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(match).flatten()\n",
    "lx = masks_pred.shape[1]\n",
    "ly = masks_pred.shape[0]\n",
    "# (tr[1,:,0]+1)*lx/2,X\n",
    "# fig,ax = plt.subplots()\n",
    "fig,ax = plt.subplots(figsize=(ly*px,lx*px))\n",
    "ax.set_aspect('equal')\n",
    "ax.axis('off')\n",
    "# ax.imshow()\n",
    "ax.set_position([0, 0, 1, 1])\n",
    "# ax.imshow(np.zeros((ly,lx,4)))\n",
    "# plt.imshow(gt_flow)\n",
    "# for i in range(tr.shape[1]):\n",
    "for i in select_inds:\n",
    "#     xs = (tr[1,i,:]+1)*lx/2\n",
    "#     ys = (tr[0,i,:]+1)*ly/2\n",
    "    xs = tr[1,i,:]\n",
    "    ys = tr[0,i,:]\n",
    "    ax.plot(xs,ys,c=[0,0,0,.5],solid_capstyle='round',linewidth=np.sqrt(A))\n",
    "\n",
    "for i in range(tr.shape[1]):\n",
    "    ax.plot(tr[1,i,-1],tr[0,i,-1],'r.',markersize=px*lx)\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",
    "plt.show()\n",
    "fig.savefig(basedir + '_perim_flows' + suffix + '.svg',bbox_inches=0,transparent=True,pad_inches=0)\n",
    "fig.savefig(basedir + '_perim_flows' + suffix + '.pdf',bbox_inches='tight',transparent=True,pad_inches=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "outline = plot.outline_view(rescale(phase),masks_pred,color=cmap(.85)[:3])#,mode='thick')\n",
    "plt.imshow(outline)\n",
    "skimage.io.imsave(basedir + '_outlines'+suffix+'.png',(outline*(2**8-1)).astype(np.uint8))\n",
    "# np.any(masks_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(masks_pred)\n",
    "# plt.imshow(dists)\n",
    "# plt.axis('off')\n",
    "# plt.show()\n",
    "# io.imsave(basedir + '_dists' + '.png',(rescale(dists)*255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "if omni:\n",
    "    a = 1\n",
    "else:\n",
    "    a = 1/6\n",
    "outlines = utils.masks_to_outlines(masks)\n",
    "outY, outX = np.nonzero(outlines)\n",
    "pic = np.stack([rescale(heat)]*3,axis=-1)**a\n",
    "pic[outY,outX] = np.array([1,0,0])\n",
    "if not omni:\n",
    "    y, x = np.nonzero(masks)\n",
    "    ymed = int(round(np.mean(y)))\n",
    "    xmed = int(round(np.mean(x)))\n",
    "    pic[ymed,xmed]=np.array([1,1,0])\n",
    "    imin = np.argmin((x-xmed)**2 + (y-ymed)**2) \n",
    "    xmed = np.array([x[imin]],np.int32)\n",
    "    ymed = np.array([y[imin]],np.int32)\n",
    "    pic[ymed,xmed]=np.array([0,1,1])\n",
    "\n",
    "pic = cmap(rescale(heat)**a)\n",
    "# pic[:,:,-1] = rescale(heat**a)\n",
    "pic[:,:,-1] = masks>0\n",
    "plt.imshow(pic)\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "skimage.io.imsave(basedir + '_8bit_heat_' + suffix + '.png',((pic)*255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# omni = True\n",
    "ext = '.tif'\n",
    "\n",
    "if omni:\n",
    "    suffix = '_OM'\n",
    "else:\n",
    "    suffix = '_OP'\n",
    "    \n",
    "flows = imread(basedir + '_flows' + suffix + ext)\n",
    "cellprob = imread(basedir + '_cellprob' + suffix + ext)\n",
    "pic = plot.dx_to_circ([flows[0,:,:],flows[1,:,:]])\n",
    "\n",
    "print(flows.shape)\n",
    "plt.figure(figsize=(10,10))\n",
    "# plt.imshow(plot.dx_to_circ([flows[:,:,1],flows[:,:,2]]))\n",
    "plt.imshow(pic)\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "# dP = flows[0,-2:,:,:]\n",
    "\n",
    "io.imsave(basedir + '_8bit_flow' + suffix + '.png',(pic*(2**8-1)).astype(np.uint8))"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}