scripts/Figure_S3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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",
    "\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",
    "import scipy.ndimage as ndimg\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",
    "\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 = 1\n",
    "basedir = '/home/kcutler/DataDrive/omnipose/docs/_static/'\n",
    "basedir = basedir+'ecoli_phase'\n",
    "maskdir = basedir+'_labels.tif'\n",
    "phasedir =  basedir+ '.tif'\n",
    "import fastremap\n",
    "masks = fastremap.renumber(imread(maskdir).astype(np.uint16))[0]\n",
    "phase = imread(phasedir)\n",
    "\n",
    "from cellpose_omni import models, io, dynamics, plot\n",
    "import edt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import omnipose\n",
    "import skimage.io\n",
    "# plt.imshow(masks)\n",
    "pic = omnipose.utils.rescale(np.stack([masks]*4,axis=-1))\n",
    "pic[...,-1] = (masks>0)\n",
    "skimage.io.imsave(basedir+'_masks_grayscale.png',omnipose.utils.to_8_bit(pic))\n",
    "skimage.io.imsave(basedir+'_masks_rgb.png',omnipose.utils.to_8_bit(omnipose.plot.apply_ncolor(masks)))\n",
    "\n",
    "\n",
    "from omnipose.plot import imshow\n",
    "imshow(pic)\n",
    "np.unique(pic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# np.unique(masks)\n",
    "# imshow(omnipose.utils.to_8_bit(pic))\n",
    "# omnipose.utils.to_8_bit(pic)\n",
    "omnipose.plot.apply_ncolor(masks).shape, pic.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "omni = False\n",
    "if omni:\n",
    "    suffix = '_OM'\n",
    "else:\n",
    "    suffix = '_CP'\n",
    "labels,dists,boundaries,heat,flows = omnipose.core.masks_to_flows(masks,use_gpu=0,omni=0,normalize=True)#,smooth=True)\n",
    "dy = flows[0,:,:].cpu().numpy()\n",
    "dx = flows[1,:,:].cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flows.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# pic = plot.dx_to_circ([dy,dx])\n",
    "pic = plot.dx_to_circ([dy,dx],transparency=True)\n",
    "# pic = omnipose.plot.rgb_flow(flows.unsqueeze(0),transparency=True).squeeze()\n",
    "# pic = rescale(heat)\n",
    "imshow(pic)\n",
    "# plt.imshow(np.arctan2(dy,dx))\n",
    "# plt.imshow(np.logical_or(dx==0, dy==0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# io.imsave('/home/kcutler/DataDrive/tle3TIFFs/d1/align/image_backup/frames/frame1_heatflow_grad_color.png', cv2.resize((pic*255).astype(np.uint8), dsize=(10*pic.shape[1],10*pic.shape[0]), interpolation=cv2.INTER_NEAREST))\n",
    "# io.imsave('/home/kcutler/DataDrive/cephalexin/OG_omni_flow_training.png', cv2.resize((pic*255).astype(np.uint8), dsize=(10*pic.shape[1],10*pic.shape[0]), interpolation=cv2.INTER_NEAREST))\n",
    "# io.imsave('/home/kcutler/DataDrive/neuronomniflow_original.png', cv2.resize((pic*255).astype(np.uint8), dsize=(10*pic.shape[1],10*pic.shape[0]), interpolation=cv2.INTER_NEAREST))\n",
    "skimage.io.imsave(basedir+'_GT.png',omnipose.utils.to_8_bit(pic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pic = 1-heat\n",
    "io.imsave(basedir + '_heat.png',(pic*255).astype(np.uint8))\n",
    "pic = 1-np.sqrt(heat)\n",
    "io.imsave(basedir + '_sqheat.png',(pic*255).astype(np.uint8))\n",
    "\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "pic = cmap(np.sqrt(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": [
    "pic = 1-(masks>0)*.5\n",
    "io.imsave(basedir + '_prob_GT.png',(pic*255).astype(np.uint8))\n",
    "pic = 1-rescale(masks)\n",
    "io.imsave(basedir + '_label_GT.png',(pic*255).astype(np.uint8))\n",
    "\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "pic = cmap(.95*(masks>0))\n",
    "pic[:,:,-1] = masks>0\n",
    "plt.imshow(pic)\n",
    "skimage.io.imsave(basedir + '_color_prob' + suffix + '.png',((pic)*255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# name = 'dnaAmovie_t112c1'\n",
    "# name = 'xy1crop2_t300c1'\n",
    "# name = '022'\n",
    "# name = '005'\n",
    "# name = 'f_xy6c1'\n",
    "# name = 'vibrio'\n",
    "# name = 'bthai2'\n",
    "name = 'caulo4'\n",
    "# name = 'serra'\n",
    "name = 'ec2'\n",
    "# name = 'ec111'\n",
    "# name = 'ect'\n",
    "# name = 'ect2'\n",
    "sourcedir = '/home/kcutler/DataDrive/omnipose_paper/test/'\n",
    "cellprob = imread(sourcedir + name + '_cellprob' + '.tif')\n",
    "# basedir =  '/home/kcutler/DataDrive/omnipose_paper/Figure 2/ecoli_phase_2'\n",
    "imshow(cellprob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from omnipose.utils import rescale\n",
    "pic = 1-rescale(cellprob)*.5\n",
    "io.imsave(basedir + '_prob_rgb.png',np.stack([pic*255]*3,axis=-1).astype(np.uint8))\n",
    "io.imsave(basedir + '_prob-raw.tif',pic)\n",
    "pic = 1-(cellprob>-1)*.25\n",
    "io.imsave(basedir + '_prob_cutoff_rgb.png',np.stack([pic*255]*3,axis=-1).astype(np.uint8))\n",
    "\n",
    "\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "pic = cmap(.95*rescale(cellprob))\n",
    "pic[:,:,-1] = rescale(cellprob)\n",
    "plt.imshow(pic)\n",
    "skimage.io.imsave(basedir + '_color_prob_predicted' + suffix + '.png',((pic)*255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flows = imread(sourcedir + name + '_flows' + '.tif')\n",
    "print(flows.shape)\n",
    "plt.figure(figsize=(10,10))\n",
    "dP = np.stack([flows[0,:,:],flows[1,:,:]])\n",
    "plt.imshow(plot.dx_to_circ(dP,transparency=True))\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pic = plot.dx_to_circ([flows[0,:,:],flows[1,:,:]],transparency=True)\n",
    "# skimage.io.imsave(basedir + '_flow_big.png', cv2.resize((pic*255).astype(np.uint8), dsize=(10*pic.shape[1],10*pic.shape[0]), interpolation=cv2.INTER_NEAREST))\n",
    "skimage.io.imsave(basedir + '_flow.png',(pic*255).astype(np.uint8))\n",
    "#replace cellpose cv2 with skimage imsave for transparency \n",
    "\n",
    "io.imsave(basedir + '_Dx.tif',flows[1,:,:])\n",
    "io.imsave(basedir + '_Dy.tif',flows[0,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "masks = imread(sourcedir+name+'_cp_masks_OG.png')\n",
    "pic = omnipose.utils.rescale(np.stack([masks]*4,axis=-1))\n",
    "pic[...,-1] = (masks>0)\n",
    "\n",
    "skimage.io.imsave(basedir+'_cp_masks_OG_greyscale.png',omnipose.utils.to_8_bit(pic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# pic = omnipose.utils.rescale(np.stack([masks>0]*4,axis=-1))/2\n",
    "pic = cmap((masks>0)*.95)\n",
    "pic[...,-1] = masks>0\n",
    "imshow(pic)\n",
    "skimage.io.imsave(basedir+'_thresh.png',omnipose.utils.to_8_bit(pic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omni=0\n",
    "if omni:\n",
    "    suffix = ''\n",
    "else:\n",
    "    suffix = '_OG'\n",
    "\n",
    "# from cellpose_omni import utils\n",
    "masks, p, tr = dynamics.compute_masks(dP, cellprob, None, p=None, niter=200, mask_threshold=-1, diam_threshold=12.,\n",
    "                                      flow_threshold=0.4, interp=True, cluster=False, do_3D=False, \n",
    "                                      min_size=15, resize=None, omni=omni, calc_trace=1, verbose=1)\n",
    "np.any(masks)"
   ]
  },
  {
   "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",
    "dists = edt.edt(masks)\n",
    "bd = dists==1\n",
    "Y,X = np.nonzero(bd) #boundary only\n",
    "# select_inds = np.flatnonzero(masks)\n",
    "a = 1\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 = cellprob.shape[1]\n",
    "ly = cellprob.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.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,.25],solid_capstyle='round',linewidth=np.sqrt(A))\n",
    "\n",
    "\n",
    "for i in range(tr.shape[1]):\n",
    "    ax.plot(tr[1,i,-1],tr[0,i,-1],'r.',markersize=2*np.sqrt(A)*px*lx)\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",
    "phasedir =  basedir + '.tif'\n",
    "phase = imread(phasedir)\n",
    "outline = plot.outline_view(phase/255,masks,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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# flow legend\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect('equal')\n",
    "ax.axis('off')\n",
    "rs = [[.1,.4],[.45,.95],[1,2]]\n",
    "l = .19\n",
    "n = [7,14,28]\n",
    "d = [0,0,0]\n",
    "for k in range(len(rs)):\n",
    "    r1,r2 = rs[k]\n",
    "    for theta in np.linspace(0,2*np.pi*(n[k]-1)/n[k],n[k]):\n",
    "        theta+=d[k]\n",
    "#         print(theta)\n",
    "        r = ((np.cos(theta)+1)/2)\n",
    "        g = ((np.cos(theta+2*np.pi/3)+1)/2)\n",
    "        b = ((np.cos(theta+4*np.pi/3)+1)/2)\n",
    "        \n",
    "        theta =theta+np.pi/2\n",
    "        ct = np.cos(theta)\n",
    "        st = np.sin(theta)\n",
    "        w = (r2-r1)\n",
    "        plt.arrow(r1*ct,r1*st,(r2-r1)*ct,(r2-r1)*st,head_width=l,\n",
    "                  head_length=l,length_includes_head=1,facecolor=[r,g,b,(r2-r1)], \n",
    "                  width=.1,edgecolor=[0,0,0,0])\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_paper/Figure images/'\n",
    "fig.savefig(os.path.join(basedir,'flow_legend.pdf'),bbox_inches='tight',transparent=True,pad_inches=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}