scripts/Figures_3_4_S1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Masks to ground truth labels: Omnipose vs Cellpose\n",
    "This script pulls out information needed to illustrate the difference in operaton between Cellpose and Omnipose. First, in Extended Data Figure 1, we plot some size and shape statistics for cells in the test and train set to confirm that there is no difference and possible bias in performance benchmarks. \n",
    "\n",
    "Next, we look just at how ground truth data is processed for training the network and how cell centers are not suitable for arbitrary morphologies. This constitutes the Extended Data Figures 4 and 5. \n",
    "\n",
    "Finally, we select some morphologically interesting cells for Figure 2 to demonstrate how the Omnipose prediction classes generalize to arbitrary shapes and how masks are reconstructed in idealized conditions (i.e. using a ground-truth flow field). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "\n",
    "from pathlib import Path\n",
    "from cellpose_omni import plot, models, core, utils, io, dynamics, transforms\n",
    "import omnipose\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": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Define datasets and directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'cyto2'\n",
    "train_dir = '/home/kcutler/DataDrive/cyto2/train'\n",
    "test_dir = '/home/kcutler/DataDrive/cyto2/test'\n",
    "img_filter = '_img'\n",
    "\n",
    "dataset = 'bact_phase'\n",
    "train_dir = '/home/kcutler/DataDrive/omnipose_all/phase/train_sorted'\n",
    "test_dir = '/home/kcutler/DataDrive/omnipose_all/phase/test_sorted'\n",
    "img_filter = ''\n",
    "\n",
    "basedir = [train_dir,test_dir]\n",
    "subset = ['train','test']\n",
    "\n",
    "save0 = os.path.join('/home/kcutler/DataDrive/omnipose_all',dataset+'_cell_metrics')\n",
    "io.check_dir(save0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Load ground-truth images and labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_GT(basedir):\n",
    "    mask_filter = '_masks'\n",
    "    img_names = io.get_image_files(basedir, mask_filter, img_filter=img_filter, look_one_level_down=True)\n",
    "    mask_names = io.get_label_files(img_names, mask_filter, img_filter=img_filter)\n",
    "    #annoyingly, the names are not sorted properly because the paths are all different for the test folder...\n",
    "    def getname(path,suffix='_masks'):\n",
    "        return os.path.splitext(Path(path).name)[0].replace(suffix,'')\n",
    "    mask_names = sorted(mask_names,key=getname)\n",
    "    img_names = sorted(img_names,key=getname)\n",
    "    imgs = [transforms.move_min_dim(skimage.io.imread(f)) for f in img_names]\n",
    "    # masks_gt = [utils.format_labels(utils.clean_boundary(skimage.io.imread(f))) for f in mask_names]\n",
    "    masks_gt = [omnipose.utils.format_labels(omnipose.utils.clean_boundary(skimage.io.imread(f))) for f in mask_names]\n",
    "\n",
    "    return imgs, masks_gt, img_names, mask_names\n",
    "def f(k):\n",
    "    if k == 0:\n",
    "        return 1,2,3,4\n",
    "    else:\n",
    "        return 4,3,2,1\n",
    "imgs =  [[]]*2\n",
    "masks_gt =  [[]]*2\n",
    "img_names =  [[]]*2\n",
    "mask_names =  [[]]*2\n",
    "for k in range(2):\n",
    "    imgs[k], masks_gt[k], img_names[k], mask_names[k] = get_GT(basedir[k])\n",
    "# train_imgs, train_masks, train_img_names, train_mask_names = get_GT(train_dir)\n",
    "# test_imgs, test_masks, test_img_names, train_mask_names = get_GT(train_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# count images to verify they are all there \n",
    "x = 0\n",
    "for k in range(2):\n",
    "    c = len(imgs[k])\n",
    "    print(c)\n",
    "    x+=c\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Data collection\n",
    "We loop over all the images in the test and train sets. For each set, we aggregate cell area, centroid, and whether or not that centroid gets projected. We copy the cell data into good and bad lists depending on this projection status. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.ndimage\n",
    "from skimage import measure\n",
    "import edt\n",
    "import cv2\n",
    "def getname(path,suffix='.tif'):\n",
    "        return os.path.splitext(Path(path).name)[0].replace(suffix,'')\n",
    "\n",
    "\n",
    "pad = 10\n",
    "\n",
    "bad_cell = [[] for k in range(2)]\n",
    "good_cell = [[] for k in range(2)]\n",
    "all_cells = [[] for k in range(2)]\n",
    "all_regs = [[] for k in range(2)]\n",
    "d2b = [[] for k in range(2)]\n",
    "md2b = [[] for k in range(2)]\n",
    "ang = [[] for k in range(2)]\n",
    "error_count = 0\n",
    "# good_or_bad =[[] for k in range(2)]\n",
    "for k in range(2):\n",
    "    for j,masks in enumerate(masks_gt[k]):\n",
    "        name = getname(img_names[k][j])\n",
    "        # Get the dimensions of the mask, preallocate arrays to store flow values\n",
    "        Ly, Lx = masks.shape\n",
    "        mu = np.zeros((2, Ly, Lx), np.float64)\n",
    "        mu_c = np.zeros((Ly, Lx), np.float64)\n",
    "\n",
    "        regions = measure.regionprops(masks)\n",
    "        all_regs[k].append(regions)\n",
    "        for i,rg in enumerate(regions):\n",
    "            label = rg.label\n",
    "            bbx = rg.bbox\n",
    "            y0 = max(0,bbx[0]-pad)\n",
    "            y1 = min(Ly,bbx[2]+pad)\n",
    "            x0 = max(0,bbx[1]-pad)\n",
    "            x1 = min(Lx,bbx[3]+pad)\n",
    "            labels = masks[y0:y1,x0:x1] \n",
    "            mask = labels==label\n",
    "            if rg.area>5:\n",
    "                # x, y ordered list of componenets for the mask pixels\n",
    "                y,x = np.nonzero(mask)\n",
    "                contours = cv2.findContours((mask>0).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)\n",
    "                x_,y_ = np.concatenate(contours[-2], axis=0).squeeze().T \n",
    "\n",
    "                # original boundary projection\n",
    "                ymed = np.median(y)\n",
    "                xmed = np.median(x)\n",
    "                imin = np.argmin((x-xmed)**2 + (y-ymed)**2) \n",
    "                xmed_ = x[imin]\n",
    "                ymed_ = y[imin]\n",
    "\n",
    "                dist_to_bound = np.sqrt((x_-xmed)**2 + (y_-ymed)**2) #absolute distance to boundary\n",
    "                imin = np.argmin(dist_to_bound)\n",
    "                min_dist_to_bound =dist_to_bound[imin]\n",
    "                l = len(dist_to_bound)\n",
    "                diff_ind = int(l/2-imin) % l\n",
    "                \n",
    "                # roll the minimum to the center\n",
    "                dist_to_bound = np.roll(dist_to_bound,diff_ind)\n",
    "                x_ = np.roll(x_,diff_ind)\n",
    "                y_ = np.roll(y_,diff_ind)\n",
    "                \n",
    "                angles = np.array([np.arctan2([y_[i]-ymed],[x_[i]-xmed]) for i in range(len(y_))])\n",
    "                \n",
    "                image = imgs[k][j][y0:y1,x0:x1]\n",
    "                data = [mask,image,dist_to_bound,angles,ymed,xmed,ymed_,xmed_,y_,x_,name,labels]\n",
    "                if np.abs(ymed-ymed_)>.5 or np.abs(xmed-xmed_)>.5: # if projected \n",
    "                    min_dist_to_bound *= -1.\n",
    "                    bad_cell[k].append(data)\n",
    "                else:\n",
    "                    good_cell[k].append(data)\n",
    "                    \n",
    "                all_cells[k].append(data)\n",
    "                \n",
    "                md2b[k].append(min_dist_to_bound)\n",
    "                d2b[k].append(dist_to_bound)\n",
    "                ang[k].append(angles)\n",
    "            else:\n",
    "                # in 47k labels, I admit there are 10 instances of 1-5px labels that are probably mistakes\n",
    "                # These are automatically removed during training, and will minimally impact evaluation\n",
    "                print('Size error',i,j,k,label,rg.area) \n",
    "                error_count += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Consolidate single-cell data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "areas = [[rg.area for region in all_regs[k] for rg in region] for k in range(2)]\n",
    "diameters = [[utils.diameters(np.pad(rg.image,1),omni=True)[0] for region in all_regs[k] for rg in region] for k in range(2)]\n",
    "perimeters = [[rg.perimeter for region in all_regs[k] for rg in region] for k in range(2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the ratio of the mean distance to boundary to the mean diameter.\n",
    "ratio = [[i / j for i, j in zip(md2b[k],diameters[k])] for k in range(2)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Calculate percentage of projected cells and cells with with 'off center' center\n",
    "Here 'off center' is defined as the mean distance to boundary being less than 20% of the mean diameter. The mean distance to boundary should be equal to .5 for a 'centered' point, at least in the circular case. The following writes this information to a text file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = ['Total cells: {}'.format(np.sum([len(ac) for ac in all_cells]))+'\\n']\n",
    "bad_percent = [[]]*2\n",
    "good_percent = [[]]*2\n",
    "\n",
    "for k in range(2):\n",
    "    nbad = len(bad_cell[k])\n",
    "    ngood = len(good_cell[k])\n",
    "    total1 = nbad+ngood\n",
    "    total2 = len(all_cells[k])\n",
    "    if total1==total2:\n",
    "        total=total1\n",
    "    else:\n",
    "        print('error',total1,total2)\n",
    "    bad_percent[k] = nbad/total*100\n",
    "    good_percent[k] = ngood/total*100\n",
    "    \n",
    "    \n",
    "    total = len(ratio[k])\n",
    "    new_nbad = np.count_nonzero(np.array(ratio[k])<.2)\n",
    "    new_bad_percent = new_nbad/total*100\n",
    "    text += [subset[k],'----------',\n",
    "             'Cell count = {} \\nFraction with/without projection: {}, {}'.format(total,bad_percent[k],good_percent[k]),\n",
    "             'Number/Percent of cells with projected center: {} / {}'.format(nbad,bad_percent[k]),\n",
    "             'Number/Percent of cells with 20% off-center centers: {} / {}'.format(new_nbad,new_bad_percent),\n",
    "             'Difference: {}'.format(new_bad_percent  - bad_percent[k]),\n",
    "             '\\n']\n",
    "        \n",
    "with open(os.path.join(save0,dataset +'_centroid_metrics.txt'), \"w\") as text_file:\n",
    "    for t in text:\n",
    "        [print(t, file=f) for f in [text_file,None]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Extended Data Figure 1\n",
    "Explore cell properties to confirm that test dataset is representative of training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "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",
    "%matplotlib inline\n",
    "darkmode = False\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(n)]\n",
    "    background_color = 'k'\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",
    "    \n",
    "mpl.rcParams['figure.dpi'] = 300\n",
    "\n",
    "golden = (1 + 5 ** 0.5) / 2\n",
    "sz = 2\n",
    "labelsize = 7\n",
    "axcol = 'k'\n",
    "colors = ['r',[.5]*3]\n",
    "\n",
    "\n",
    "labels = ['Center to boundary \\ndistance (px)','Mean diameter (px)',\n",
    "          'Center-to-boundary / diameter','Area (px$^2$)','Perimeter (px)']\n",
    "data = [md2b,diameters,ratio,areas,perimeters]\n",
    "J = len(data)\n",
    "\n",
    "fig,axs = plt.subplots(1,J,figsize=(sz*J,sz))\n",
    "D = .75\n",
    "plt.subplots_adjust(wspace=D, hspace = D)\n",
    "\n",
    "# print some results as they are calculated and plotted \n",
    "with open(os.path.join(save0,dataset +'_ks_test.txt'), \"w\") as text_file:\n",
    "    for j in range(J):\n",
    "    # for j in [1,3,4]:\n",
    "        print(labels[j],file=text_file)\n",
    "        print(scipy.stats.ks_2samp(data[j][0],data[j][1]),file=text_file)\n",
    "        ax = axs[j]\n",
    "        x = [-.25,.25]\n",
    "        parts = ax.violinplot(data[j],x,showmeans=0, showextrema=0,widths=.4,points=1000)\n",
    "        for i,pc in enumerate(parts['bodies']):\n",
    "            pc.set_facecolor(colors[i])\n",
    "            pc.set_alpha(1)\n",
    "        #     pc.set_edgecolor('none')\n",
    "    #         pc.set_edgecolor(colors[i])\n",
    "\n",
    "        ax.set_xticks(x)\n",
    "        ax.set_xticklabels(['test','train'])#,rotation = 45)\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.set_ylabel(labels[j], fontsize = labelsize)\n",
    "    #     if j != 1:\n",
    "    #         ax.set_yscale('log')\n",
    "    #     if j == 3:\n",
    "    #         ax.set_ylim(0,1000)\n",
    "    #     if j == 4:\n",
    "    #         ax.set_ylim(0,1000)\n",
    "        m = np.mean(data[j][0])\n",
    "        s = np.std(data[j][0])\n",
    "        low = max(m-2*s,np.min(data[j][0]))\n",
    "        high = min(m+2*s,np.max(data[j][0]))\n",
    "        ax.set_ylim(low,high)\n",
    "        ax.spines['right'].set_visible(False)\n",
    "        ax.spines['top'].set_visible(False)\n",
    "\n",
    "    # plt.yscale('log')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Save the figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig.savefig(os.path.join(save0,'cell_statistics_'+dataset+'.eps'),bbox_inches=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Extended Data Figures 4 and 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = ratio[0]\n",
    "m = np.mean(r)\n",
    "s = np.std(r)\n",
    "t = len(r)\n",
    "np.sum(r<m)/t*100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dists = np.array([np.abs(bad_cell[0][i][2]) for i in range(len(bad_cell[0]))])\n",
    "# np.argmax(dists)\n",
    "# len(bad_cell[0])\n",
    "dists = [np.abs(bad_cell[0][i][2]) for i in range(len(bad_cell[0]))]\n",
    "angles = [bad_cell[0][i][3] for i in range(len(bad_cell[0]))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = 100\n",
    "\n",
    "th_bad = []\n",
    "r_bad = []\n",
    "for i in range(len(bad_cell[0])):\n",
    "    r = np.abs(bad_cell[0][i][2]).flatten()\n",
    "    theta = bad_cell[0][i][3].flatten()\n",
    "#     selection = r<=np.percentile(r,p)\n",
    "#     select_theta = theta[selection]\n",
    "    th_bad.append(theta)\n",
    "    r_bad.append(r)\n",
    "th_good = []\n",
    "r_good = []\n",
    "for i in range(len(good_cell[0])):\n",
    "    r = np.abs(good_cell[0][i][2]).flatten()\n",
    "    theta = good_cell[0][i][3].flatten()\n",
    "#     selection = r<=np.percentile(r,p)\n",
    "#     select_theta = theta[selection]\n",
    "    th_good.append(theta)\n",
    "    r_good.append(r)\n",
    "th_good = np.array([item for sublist in th_good for item in sublist])\n",
    "th_bad  = np.array([item for sublist in th_bad for item in sublist])\n",
    "r_good = np.array([item for sublist in r_good for item in sublist])\n",
    "r_bad  = np.array([item for sublist in r_bad for item in sublist])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 'all'\n",
    "if t=='bad':\n",
    "    group = bad_cell\n",
    "    suffix = '_bad_cell'\n",
    "elif t=='good':\n",
    "    group = good_cell\n",
    "    suffix = '_good_cell'\n",
    "elif t=='all':\n",
    "    group = all_cells\n",
    "    suffix = '_all_cells'   \n",
    "nr = [np.abs(group[0][i][2]) for i in range(len(group[0]))]\n",
    "th = [np.abs(group[0][i][3]) for i in range(len(group[0]))]\n",
    "\n",
    "X = []\n",
    "Y = []\n",
    "for k in range(len(nr)):\n",
    "    y = nr[k]/np.max(nr[k])\n",
    "    imin = np.argmin(y)\n",
    "    x = np.linspace(-1,1,len(nr[k]))#-imin/len(y)\n",
    "#     sind = np.argsort(y) # scatter plot sorted by y or x but re-sorting both y or x gives same plot, just appear in different order \n",
    "    # really what I want is to make sure that the x coordinate is unbiased... that is, show how distances change sa you go further away from the projected point\n",
    "    #right now it is basically that, but each is shifted \n",
    "    \n",
    "    \n",
    "    #INDEXING is critical... right now assumes that the x a xis spreads out cells uniformly with distance \n",
    "    \n",
    "    # plotting y vs theta (y=r) recapitulates cell shapes, sort of\n",
    "#     theta = np.concatenate(th[k])\n",
    "#     sind = np.argsort(theta)\n",
    "#     thmin = theta[imin]\n",
    "#     theta = theta - thmin\n",
    "#     theta = theta[sind]\n",
    "\n",
    "#     x = x[sind]\n",
    "#     y = y[sind]\n",
    "    Y.append(y)\n",
    "#     X.append(theta)\n",
    "    X.append(x)\n",
    "    \n",
    "X = np.array([item for sublist in X for item in sublist])\n",
    "Y = np.array([item for sublist in Y for item in sublist])\n",
    "\n",
    "\n",
    "# Construct 2D histogram from data using the 'plasma' colormap\n",
    "from scipy.ndimage import gaussian_filter\n",
    "heatmap, xedges, yedges = np.histogram2d(Y, X, bins=(300,600))\n",
    "# heatmap = gaussian_filter(heatmap, sigma=1)\n",
    "extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]\n",
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "v = omnipose.utils.rescale(np.log(np.flipud(heatmap)+1))\n",
    "pic = cmap(v)\n",
    "pic[:,:,-1] = v>0\n",
    "\n",
    "#save just the image\n",
    "# basedir = '/home/kcutler/DataDrive/Omnipose_paper/Figure 3/'\n",
    "\n",
    "name = '_projected_normalized_distance'+suffix\n",
    "filebase = os.path.join(save0, name + '_' + dataset)\n",
    "\n",
    "skimage.io.imsave(filebase +'_just_2DH_hist' + '.png',((pic)*255).astype(np.uint8))\n",
    "plt.imshow(pic)\n",
    "plt.show()\n",
    "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':['Arial']})\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 = .75\n",
    "labelsize = 8\n",
    "px = A/plt.rcParams['figure.dpi']  # pixel in inches\n",
    "lx = pic.shape[1]\n",
    "ly = pic.shape[0]\n",
    "fig,ax = plt.subplots(figsize=(ly*px,lx*px))\n",
    "\n",
    "\n",
    "plt.imshow(pic, extent=[-1,1,0,1], origin='upper',interpolation='nearest')\n",
    "plt.xlim(-1,1)\n",
    "ax = plt.gca()\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.set_ylabel('Relative center-to-boundary \\ndistance',size=labelsize)\n",
    "ax.set_xlabel('Relative position along boundary',size=labelsize)\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_aspect(1/2)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "fig.savefig(filebase + '.svg',bbox_inches=0,transparent=True,pad_inches=0)\n",
    "fig.savefig(filebase + '.pdf',bbox_inches='tight',transparent=True,pad_inches=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "### select range of images to inspect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "k = 0\n",
    "_names = [group[k][i][-2] for i in range(len(group[k]))]\n",
    "# print(_names)\n",
    "# name = 'Hpylori'\n",
    "name = 'Serratia'\n",
    "name = 'wiggins'\n",
    "name = 'streptomyces'\n",
    "# name = 'bthai'\n",
    "# name = '5I'\n",
    "# indices = [i for i, s in enumerate(img_names[0]) if name in getname(s,suffix='.tif')]\n",
    "indices = [i for i, s in enumerate(_names) if name in s]\n",
    "\n",
    "# indices = range(len(_names))\n",
    "# np.unique(_names)\n",
    "# [_names[i] for i in indices]\n",
    "len(indices),len(group)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import rc\n",
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42\n",
    "mpl.rcParams['font.family'] = 'sans-serif'\n",
    "mpl.rcParams['font.sans-serif'] = 'Arial'\n",
    "# rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})\n",
    "\n",
    "import matplotlib_inline\n",
    "# matplotlib_inline.backend_inline.set_matplotlib_formats('retina', 'png')\n",
    "mpl.rcParams['figure.dpi'] = 72\n",
    "\n",
    "\n",
    "from scipy.signal import argrelextrema\n",
    "from math import gcd\n",
    "from scipy.ndimage import binary_erosion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "savedir = os.path.join(save0,'examples')\n",
    "io.check_dir(savedir)\n",
    "save = 1\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "# group = all_cells\n",
    "show_im = 1\n",
    "plot_closest = 0\n",
    "group = all_cells\n",
    "# for ind in range(len(group[0])):\n",
    "np.random.seed(42)\n",
    "bg = 0.5\n",
    "# for ind in np.random.choice(indices,12):\n",
    "# for ind in indices:\n",
    "# for ind in [12304, 11870 ,12193, 12228, 12284, 12313]: #Hpylori\n",
    "# for ind in [23295]: # e coli wt\n",
    "# for ind in [8259,3676,7665,3732,23295]: #all\n",
    "# for ind in [11870]:\n",
    "# for ind in [34,96,116]:\n",
    "for ind in indices:\n",
    "    r = group[0][ind][2]\n",
    "    theta = np.concatenate(group[0][ind][3])\n",
    "#     inds = np.argsort(theta)\n",
    "#     r = r[inds]\n",
    "\n",
    "    # [mask,phase,dist_to_bound,angles,ymed,xmed,ymed_,xmed_,y_,x_,name,labels\n",
    "#     if nr<.2:\n",
    "    argmins = argrelextrema(r, np.less,mode='wrap',order=3)[0]\n",
    "#     print('number of minma',len(argmins),argmins)\n",
    "#     if len(argmins)>4:\n",
    "#     if 1:\n",
    "    mask = np.array(group[0][ind][0]).astype(int)\n",
    "    nr = np.min(np.abs(r))/utils.diameters(mask,omni=True)[0]\n",
    "    lbl = np.array(group[0][ind][-1]).copy()\n",
    "    if len(np.unique(lbl))<3 and np.count_nonzero(mask)>200:# and nr<.2:\n",
    "        print(_names[ind])\n",
    "        p = group[0][ind][1]\n",
    "        print(p.shape)\n",
    "        # phase = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mask==0)])))\n",
    "        if dataset!='cyto2':\n",
    "            phase = omnipose.utils.normalize_image(group[0][ind][1],mask,bg)\n",
    "        else:\n",
    "            phase = group[0][ind][1]\n",
    "            \n",
    "        phase = plot.image_to_rgb(phase)/255\n",
    "\n",
    "        regions = skimage.measure.regionprops(mask)\n",
    "        coords = group[0][ind][4:-2]\n",
    "        # print('dist to boundary',bad_cell[0][ind][2])\n",
    "        print('index',ind,mask.shape)\n",
    "\n",
    "        A = .5\n",
    "        labelsize = 8\n",
    "        dpi = 72\n",
    "        plt.rcParams['figure.dpi'] = dpi\n",
    "        px = A/dpi  # pixel in inches\n",
    "        ly,lx = pic.shape[:2]\n",
    "        fig,ax = plt.subplots(figsize=(A*lx/dpi,A*ly/dpi),dpi=dpi)\n",
    "\n",
    "        # plt.xlim(-1,1)\n",
    "        ax = plt.gca()\n",
    "        ax.set_aspect('equal')\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.set_ylabel('Relative center- \\nto-boundary distance',size=labelsize)\n",
    "        ax.set_xlabel('Relative position along boundary',size=labelsize)\n",
    "        \n",
    "        \n",
    "        h,w= pic.shape[:2]\n",
    "        x = np.linspace(-1,1,len(r))#-np.argmin(r)/len(r)\n",
    "        y = r/np.max(r)\n",
    "\n",
    "\n",
    "        ax.plot(x,y,'k-',linewidth=3)\n",
    "        \n",
    "        ax.set_xlim([-1,1])\n",
    "        ax.set_ylim([0,1])\n",
    "        plt.imshow(pic, extent=[-1, 1.0, 0, 1.0], origin='upper',interpolation='none')\n",
    "\n",
    "\n",
    "        plt.show()\n",
    "    \n",
    "    \n",
    "        name = group[0][ind][-2]+'_'+str(ind)\n",
    "        if save:\n",
    "            fig.savefig(os.path.join(savedir, name  + '.svg'),bbox_inches=0,transparent=True,pad_inches=0)\n",
    "            fig.savefig(os.path.join(savedir, name  + '.eps'),bbox_inches='tight')\n",
    "            fig.savefig(os.path.join(savedir, name  + '.pdf'),bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "        \n",
    "        if show_im:\n",
    "            ly,lx = mask.shape\n",
    "            # fig,ax = plt.subplots(figsize=(lx/dpi,ly/dpi), dpi=dpi)\n",
    "            fig,ax = plt.subplots(figsize=(10,10),dpi=dpi)\n",
    "            \n",
    "            labels,dists,heat,flows = omnipose.core.masks_to_flows(mask,use_gpu=True,omni=False)\n",
    "            dP = np.stack([flows[0],flows[1]])\n",
    "            ol = plot.outline_view(phase,mask,color=cmap(.85)[:3])\n",
    "            cmap = mpl.cm.get_cmap('plasma')\n",
    "            pic2 = cmap(omnipose.utils.rescale(heat))\n",
    "            pic2[:,:,-1] = mask>0\n",
    "\n",
    "            pic3 = plot.dx_to_circ(dP,transparency=True)\n",
    "\n",
    "            div = omnipose.utils.rescale(divergence(dP))\n",
    "            pic4 = cmap(div)\n",
    "            pic4[:,:,-1] = mask>0\n",
    "            if save:\n",
    "                skimage.io.imsave(os.path.join(savedir, name +'_8bit_heat' + '.png'),((pic2)*255).astype(np.uint8))\n",
    "                skimage.io.imsave(os.path.join(savedir, name +'_GT_flow'+'.png'),(pic3*(2**8-1)).astype(np.uint8))\n",
    "                skimage.io.imsave(os.path.join(savedir, name +'_GT_divergence'+'.png'),(pic4*(2**8-1)).astype(np.uint8))\n",
    "\n",
    "            \n",
    "            ax.imshow(ol,interpolation='none')\n",
    "            xmed,ymed = coords[1],coords[0]\n",
    "            x_,y_ = coords[-1],coords[-2]\n",
    "            \n",
    "            if plot_closest:\n",
    "                ax.plot(x_[argmins],y_[argmins],'.',color='mediumturquoise',markersize=5, markeredgewidth=0) #minima\n",
    "            ax.plot(xmed,ymed,'x',color=cmap(.35),markersize=2.5,mew=1) # median\n",
    "#             ax.set_xlim([0,lx])\n",
    "#             ax.set_ylim([ly,0])\n",
    "            ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "            if save:\n",
    "                skimage.io.imsave(os.path.join(savedir, name+'_outlines'+'.png'),(ol).astype(np.uint8))\n",
    "                fig.savefig(os.path.join(savedir, name  + '_outline_median_px.pdf'),\n",
    "                            bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "                fig.savefig(os.path.join(savedir, name  + '_outline_median_px.eps'),\n",
    "                            format='eps',bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "\n",
    "\n",
    "        \n",
    "#         ax =axs[2]\n",
    "#         theta = group[0][ind][3]\n",
    "#         ax.plot(theta,'r.',markersize=1)\n",
    "#         ax.set_ylabel('angle',size=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",
    "\n",
    "#         ax = axs[3]\n",
    "#         ax.scatter(r,theta,c='b',s=1)\n",
    "#         ax.set_ylabel('angle vs distance',size=labelsize)\n",
    "#         # ax = fig.add_subplot(projection='polar')\n",
    "#         # ax.scatter(theta, r, c=theta, s=1, cmap='hsv', alpha=0.75)\n",
    "\n",
    "#         # selection = r<np.percentile(r,1)\n",
    "#         # select_theta = theta[selection]\n",
    "#         # print(np.min(r),np.ptp(select_theta))\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",
    "\n",
    "        plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, \n",
    "                    hspace = 0, wspace = 1)\n",
    "        plt.show()\n",
    "\n",
    "#[mask,phase,dist_toadound,angles,diam,ymed,xmed,ymed_,xmed_]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "# Figure 4\n",
    "In figure 4, we want to "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = mpl.cm.get_cmap('plasma')\n",
    "bg = 0.5\n",
    "# for i in inds:\n",
    "from cellpose_omni import transforms\n",
    "from scipy.ndimage import binary_dilation\n",
    "\n",
    "\n",
    "basedir = '/home/kcutler/DataDrive/omnipose_all/omnipose_gt_examples/'\n",
    "io.check_dir(basedir)\n",
    "# for i in [inds[-13]]:\n",
    "# for i in inds:\n",
    "for i in [4243,19151,2889,4214,4023,4030,3955]:\n",
    "    border_mask = np.zeros(masks[i].shape, dtype=bool)\n",
    "    border_mask = binary_dilation(border_mask, border_value=1, iterations=5)\n",
    "    overlap = np.count_nonzero(np.logical_and(masks[i], border_mask))\n",
    "    if overlap==0:\n",
    "        p = imgs[i]\n",
    "        mask = masks[i]\n",
    "        phase = p**(np.log(bg)/np.log(np.mean(p[binary_erosion(mask==0)])))\n",
    "        \n",
    "        print('index',i)\n",
    "        fig = plt.figure(figsize=(2,2))\n",
    "        dist = edt.edt(masks[i])\n",
    "        smooth_dist = omnipose.core.smooth_distance(masks[i],dist)\n",
    "        mu = derivative(dist,masks[i])\n",
    "        mu = transforms.normalize_field(mu)\n",
    "        mu2 = derivative(smooth_dist,masks[i])\n",
    "        mu2 = transforms.normalize_field(mu2)\n",
    "        outline = plot.outline_view(phase,mask,color=cmap(.85)[:3],mode='thick')\n",
    "        flow1 = plot.dx_to_circ(mu,transparency=True)\n",
    "        flow2 = plot.dx_to_circ(mu2,transparency=True)\n",
    "        colordist1 = cmap(omnipose.utils.rescale(dist))\n",
    "        colordist1[:,:,-1] = masks[i]>0\n",
    "        colordist2 = cmap(omnipose.utils.rescale(smooth_dist))\n",
    "        colordist2[:,:,-1] = masks[i]>0\n",
    "        diff1 = cmap(omnipose.utils.rescale(smooth_dist-dist))\n",
    "        diff1[:,:,-1] = masks[i]>0\n",
    "        diff2 = plot.dx_to_circ(mu2-mu,transparency=True)\n",
    "\n",
    "        plt.imshow(np.vstack((np.hstack((outline, colordist1[:,:,:3], flow1[:,:,:3])),\n",
    "                              np.hstack((np.ones_like(outline), colordist2[:,:,:3],flow2[:,:,:3])),\n",
    "                              np.hstack((np.ones_like(outline), diff1[:,:,:3], diff2[:,:,:3])))\n",
    "                            ))\n",
    "        plt.axis('off')\n",
    "        plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, \n",
    "                hspace = 0, wspace = 0)\n",
    "        plt.show()\n",
    "        \n",
    "        name = 'cell_'+str(i)\n",
    "        savedir = os.path.join(basedir,name)\n",
    "        skimage.io.imsave(savedir + '_outlines'+'.png',(outline*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_FMM'+'.png',(colordist1*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_smoothdist'+'.png',(colordist2*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_FMM_flow'+'.png',(flow1*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_smoothdist_flow'+'.png',(flow2*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_dist_diff'+'.png',(diff1*(2**8-1)).astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_flow_diff'+'.png',(diff2*(2**8-1)).astype(np.uint8))\n",
    "        \n",
    "        skimage.io.imsave(savedir + '_mask'+'.png',mask.astype(np.uint8))\n",
    "        skimage.io.imsave(savedir + '_phase'+'.png',(phase*255).astype(np.uint8))\n",
    "        \n",
    "        skimage.io.imsave(savedir + '_boundary'+'.png',((dist==1)*(2**8-1)).astype(np.uint8))\n",
    "\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_names[k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "1.5*25\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks = [[rg.image for region in all_regs[k] for rg in region] for k in range(2)]\n",
    "masks = masks[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(masks[4243])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "save0"
   ]
  },
  {
   "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
}