scripts/3D_omnipose.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "from cellpose_omni import io\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",
    "from cellpose_omni import dynamics, transforms\n",
    "from scipy import stats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aicsimageio.readers import CziReader\n",
    "import os\n",
    "basedir = '/home/kcutler/DataDrive/confocal'\n",
    "dataset = 'Experiment-712' #720 too dim, 733 too thick, 704 too binned,707 good! try 708 next \n",
    "dataset = 'Experiment-707' \n",
    "dataset = 'Experiment-719' #clean up edges of my dataset\n",
    "\n",
    "reader = CziReader(os.path.join(basedir,dataset+'.czi'))\n",
    "im = reader.data.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tifffile\n",
    "# im = tifffile.imread('/home/kcutler/DataDrive/3D_BBBC/BBBC027/train/0000.tif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldir = '/home/kcutler/DataDrive/confocal/train60/models/cellpose_residual_on_style_on_concatenation_off_omni_train60_2022_04_28_02_39_03.970940_epoch_3999'\n",
    "modeldir = '/home/kcutler/DataDrive/confocal/train60/models/cellpose_residual_on_style_on_concatenation_off_omni_train60_2022_04_28_02_39_03.970940_epoch_3951'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 3\n",
    "nclasses = dim+2\n",
    "nchan = 1\n",
    "omni = 1\n",
    "rescale = False\n",
    "diam_mean = 0\n",
    "from cellpose_omni import models, core\n",
    "use_GPU = core.use_gpu()\n",
    "print('>>> GPU activated? %d'%use_GPU)\n",
    "use_GPU = 0\n",
    "model = models.CellposeModel(gpu=use_GPU, pretrained_model=modeldir, net_avg=False, \n",
    "                             diam_mean=diam_mean, nclasses=nclasses, dim=dim, nchan=nchan)\n",
    "model.pretrained_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_threshold = -5 ##############\n",
    "diam_threshold = 12\n",
    "net_avg = 0\n",
    "cluster = 0\n",
    "verbose = 1\n",
    "\n",
    "tile = 0\n",
    "chans = [0,0]\n",
    "chans = None\n",
    "compute_masks = 1\n",
    "slc_crop = slice(None)\n",
    "# slc_crop = slice(100, 120, None)\n",
    "# slc_crop = slice(13, 53, None) # 733\n",
    "# slc_crop = slice(50, 100, None) # 704\n",
    "slc_crop = slice(15, 52, None) # 707\n",
    "# slc_crop = slice(30, 70, None) # 712\n",
    "# slc_crop = (Ellipsis,)+tuple([slice(0,129)]*2)\n",
    "resample=False\n",
    "rescale = None\n",
    "\n",
    "imcrop = im[slc_crop]\n",
    "# imcrop = im[100:120]\n",
    "# imcrop = im.copy()\n",
    "# import torch\n",
    "imcrop.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks, flows, styles = model.eval(imcrop,channels = chans,rescale=rescale,mask_threshold=mask_threshold,net_avg=net_avg,\n",
    "                                  transparency=True,flow_threshold=0.,omni=omni,resample=resample,verbose=verbose,\n",
    "                                  diam_threshold=diam_threshold,cluster=cluster,tile=0,\n",
    "                                  compute_masks=compute_masks)\n",
    "#     else:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import omnipose, tifffile\n",
    "mu = flows[1]\n",
    "dt = flows[2]\n",
    "p = flows[3]\n",
    "bd = flows[4]\n",
    "suffix = '_1951'\n",
    "ext = '.tif'\n",
    "tifffile.imwrite(os.path.join(basedir,dataset+'_trained_flow'+suffix+ext),mu)\n",
    "tifffile.imwrite(os.path.join(basedir,dataset+'_trained_dist'+suffix+ext),dt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "hdbscan = 1\n",
    "cluster = 0\n",
    "masks,p,_= omnipose.core.compute_masks(mu,dt,bd=bd,mask_threshold=mask_threshold,\n",
    "                                           use_gpu=True,verbose=True,cluster=cluster, \n",
    "                                           flow_threshold=0., dim=3, hdbscan=hdbscan, \n",
    "                                        debug=0, eps=3, min_size=30, nclasses=im.ndim+2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import napari, ncolor\n",
    "from omnipose.utils import sinebow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnc = ncolor.label(masks)\n",
    "N = np.max(mnc)\n",
    "viewer = napari.view_image(imcrop, name='cells')\n",
    "viewer.add_labels(mnc, name='labels',color=sinebow(N+3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tifffile\n",
    "tifffile.imwrite(os.path.join(basedir,dataset+'_trained_masks'),masks)\n",
    "tifffile.imwrite(os.path.join(basedir,dataset+'_trained_flow'),mu)\n",
    "tifffile.imwrite(os.path.join(basedir,dataset+'_trained_dist'),dt)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.add_image(flows[2], name='dist')\n",
    "viewer.add_image(flows[4], name='bd')\n",
    "viewer.add_image(flows[0], name='flow')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = flows[1]\n",
    "mag = np.sqrt(np.nansum(mu**2,axis=0))\n",
    "viewer.add_image(mag>2.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt = ncolor.label(tifffile.imread('/home/kcutler/DataDrive/3D_BBBC/BBBC027/train/0000_masks.tif'))[slc_crop]\n",
    "N = gt.max()\n",
    "viewer.add_labels(gt, name='labels_gt',color=sinebow(N+3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "mag.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "im.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(im[:,300])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tifffile\n",
    "stack = tifffile.imread('/home/kcutler/DataDrive/luca/data/tif_stacks/191.tif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = napari.view_image(stack, name='cells')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aicsimageio import AICSImage\n",
    "idir = '/home/kcutler/DataDrive/3D_BBBC/BBBC024_v1_c75_highSNR_foreground/image-labels_0015.ics'\n",
    "img = AICSImage(idir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# img.data  # returns 5D TCZYX numpy array\n",
    "# img.xarray_data  # returns 5D TCZYX xarray data array backed by numpy\n",
    "# img.dims  # returns a Dimensions object\n",
    "# img.dims.order  # returns string \"TCZYX\"\n",
    "# img.dims.X  # returns size of X dimension\n",
    "# img.shape  # returns tuple of dimension sizes in TCZYX order\n",
    "# img.get_image_data(\"CZYX\", T=0)  # returns 4D CZYX numpy array\n",
    "import napari\n",
    "viewer = napari.view_image(img.data.squeeze(), name='cells')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install aicsimageio[bfio] --force --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tifffile\n",
    "from cloudvolume import CloudVolume\n",
    "\n",
    "# vol = CloudVolume(\n",
    "#     \"s3://open-neurodata/hildebrand/120809zf104/151009\", mip=0, use_https=True\n",
    "# )\n",
    "\n",
    "# # load data into numpy array\n",
    "# cutout = vol[1024:1536, 512:1024, 7520:7536]\n",
    "\n",
    "\n",
    "vol = CloudVolume(\n",
    "    \"s3://open-neurodata/cardona/acardona_abd1_5/image\", mip=0, use_https=True\n",
    ")\n",
    "\n",
    "# load data into numpy array\n",
    "cutout = vol[10752:11264, 9216:9728, 96:112]\n",
    "viewer = napari.view_image(cutout.squeeze(), name='cells')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = tifffile.imread('/home/kcutler/DataDrive/Pedicel/Pedicel Segmentation.tif')\n",
    "im = tifffile.imread('/home/kcutler/DataDrive/Pedicel/Pedicel stack.tif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import napari\n",
    "# mnc = ncolor.label(mask)\n",
    "# N = np.max(mnc)\n",
    "viewer = napari.view_image(im, name='cells')\n",
    "viewer.add_labels(np.flip(mask,axis=0), name='labels')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "cutout.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ncolor\n",
    "mnc0 = ncolor.label(mask,max_depth = 11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask.shape,im.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.add_labels(np.flip(mask,axis=2), name='labels2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.add_labels(np.flip(mask,axis=1), name='labels1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cellpose_omni import plot\n",
    "fl = plot.dx_to_circ(flows[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cellpose_omni import io\n",
    "save0 = os.path.join('/home/kcutler/DataDrive/confocal',dataset,'flow')\n",
    "io.check_dir(save0)\n",
    "ext = '.png'\n",
    "for k in range(im.shape[1]):\n",
    "# for k in [0]:\n",
    "    # imk = np.uint8(flows[0][:,k])\n",
    "    imk = np.uint8(fl[:,k])\n",
    "    skimage.io.imsave(os.path.join(save0,str(k)+ext),imk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "flows[1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(fl[:,4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "2755/10*40/60/60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "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": 5
}