Unified_Tracking_and_Shape_Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5mvQa2M5Ausc",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Mount Google Drive (Colab can see Drive files) and authenticate so that we can interact with GCP via SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_colab = True;\n",
    "try:\n",
    "  from google.colab import drive\n",
    "  mount_location = '/content/drive'\n",
    "  drive.mount(mount_location,force_remount=True);\n",
    "  \n",
    "  from pathlib import Path\n",
    "  import sys\n",
    "  libraries_path =  \"MyDrive\" #@param {type:\"string\"}\n",
    "  colab_libraries_path = Path(mount_location)/libraries_path\n",
    "  print(\"c_libraries_path:\",colab_libraries_path)\n",
    "  sys.path.insert(0,str(colab_libraries_path)); #add google drive's libraries to path\n",
    "  \n",
    "  from google.colab import auth\n",
    "  #This allows SDK to see and edit Google Drive files\n",
    "  #SDK is required to interact with GCP\n",
    "  auth.authenticate_user()\n",
    "except ImportError as ie:\n",
    "  print(\"google not found; assuming non-colab execution\")\n",
    "  print(ie);\n",
    "  in_colab = False;\n",
    "except Exception as e:\n",
    "  print(\"google authentication failed, please retry\")\n",
    "  raise e\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2zRp7M2Id99S",
    "tags": []
   },
   "source": [
    "# Global Parameters and Setup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lHLKohT4ht29",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## General library imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    %pip install -r requirements.txt\n",
    "except:\n",
    "    %pip install -Uqq scikit-image\n",
    "    %pip install -Uqq tifffile\n",
    "    %pip install Pillow==9.0.1\n",
    "    %pip install scikit-image==0.19.3\n",
    "    %pip install -Uqq mediapy\n",
    "\n",
    "try:\n",
    "    !command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
    "except:\n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "import numpy as np\n",
    "from ipywidgets import interact, widgets, Layout\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from skimage.io import imread, imsave, imshow\n",
    "from skimage import data, filters, measure, morphology, util\n",
    "from skimage.exposure import rescale_intensity\n",
    "from skimage.transform import rescale, resize, downscale_local_mean, rotate\n",
    "from skimage.measure import  regionprops, regionprops_table\n",
    "from pathlib import Path\n",
    "from fastprogress.fastprogress import progress_bar,master_bar\n",
    "from enum import Enum\n",
    "import tifffile\n",
    "import pandas as pd\n",
    "import cv2\n",
    "import ntpath\n",
    "import os\n",
    "import time\n",
    "import pickle\n",
    "import re\n",
    "import copy\n",
    "import scipy\n",
    "import builtins\n",
    "import shutil\n",
    "import contextlib\n",
    "import zipfile\n",
    "from itertools import starmap\n",
    "from functools import partial\n",
    "import random\n",
    "import csv\n",
    "from scipy import ndimage\n",
    "from scipy.stats import skew\n",
    "import numpy as np\n",
    "from ipywidgets import interact, widgets, Layout\n",
    "from IPython.utils.io import capture_output\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import Dict, DefaultDict, Tuple, List, Union, Iterable\n",
    "\n",
    "import itertools\n",
    "import math\n",
    "\n",
    "import mediapy as media\n",
    "\n",
    "from operator import itemgetter\n",
    "\n",
    "import libraries.filter_cells_fns\n",
    "import libraries.centers\n",
    "import libraries.qc_functions\n",
    "import libraries.centroidtracker\n",
    "import libraries.smoothing\n",
    "import libraries.parsend\n",
    "\n",
    "#Get from Drive the folder libraries\n",
    "def load_libraries():\n",
    "    import importlib\n",
    "    [importlib.reload(f) for f in [libraries.filter_cells_fns,\n",
    "                                   libraries.centers,\n",
    "                                   libraries.qc_functions,\n",
    "                                   libraries.centroidtracker,\n",
    "                                   libraries.smoothing,\n",
    "                                   libraries.parsend]]\n",
    "        \n",
    "    from libraries.filter_cells_fns import remove_multiple_nuclei_cells, remove_large_objects, remove_touching_edge\n",
    "    from libraries.centers import get_centers, fill_label_holes, normalize\n",
    "    from libraries.qc_functions import apply_qc\n",
    "    from libraries.centroidtracker import CentroidTracker\n",
    "    from libraries.smoothing import moving_average\n",
    "    from libraries.parsend import parseND, StageDict, group_stage_basenames\n",
    "    \n",
    "    # import libraries; ## THESE ARE NOT USEFUL: just to allow reloading\n",
    "    globals().update(locals())\n",
    "load_libraries()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tAx-TRKhubBH",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Auxiliary Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vNh24D32ZSTL",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### GCP Path formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_gcp_path(path:Union[str,Path]):\n",
    "  if not isinstance(path,Path):\n",
    "    path = Path(path);\n",
    "  return path.parts[0].lower() == \"gs:\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XTheZagyE0u7"
   },
   "source": [
    "### File zipping/unzipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _linux_cmd_zip(source,dest,recurse=False,compresslevel:Union[int,None]=None,relative_to=\"\"):\n",
    "  relative = source/relative_to\n",
    "  compresstext = f\"-{compresslevel}\"\n",
    "  if compresslevel is not None:\n",
    "    if recurse:\n",
    "      result = !cd \"{relative}\" && zip \"{compresstext}\" -r \"{dest.name}\" *\n",
    "    else:\n",
    "      result = !cd \"{relative}\" && zip \"{compresstext}\" \"{dest.name}\" *\n",
    "  else:\n",
    "    if recurse:\n",
    "      result = !cd \"{relative}\" && zip -r \"{dest.name}\" *\n",
    "    else:\n",
    "      result = !cd \"{relative}\" && zip \"{dest.name}\" *\n",
    "  shutil.move(relative/dest.name,dest);\n",
    "  return result\n",
    "\n",
    "def _python_cmd_zip(source,dest,recurse=False,compresslevel:Union[int,None]=6):\n",
    "  print(source);\n",
    "  if compresslevel is None:\n",
    "    compresslevel = 6\n",
    "  source = Path(source);\n",
    "  dest = Path(dest)\n",
    "  with zipfile.ZipFile(dest,'w',compression=zipfile.ZIP_DEFLATED,compresslevel=compresslevel) as archive:\n",
    "      for filepath in (source.rglob(\"*\") if recurse else source.iterdir()):\n",
    "          archive.write(filepath,arcname=filepath.relative_to(source));\n",
    "zipExists = shutil.which(\"zip\")\n",
    "cmd_zip = _linux_cmd_zip if zipExists else _python_cmd_zip\n",
    "print(\"using python zip\" if not zipExists else \"using cmdline zip\")\n",
    "\n",
    "def _linux_cmd_unzip(source,dest,overwrite=False)->List[str]:\n",
    "  flags = \"-\" + (\"o\" if overwrite else \"\")\n",
    "  if flags ==  \"-\":\n",
    "    result = !unzip \"{source}\" -d \"{dest}\" \n",
    "  else:\n",
    "    result = !unzip \"{flags}\" \"{source}\" -d \"{dest}\" \n",
    "  return result\n",
    "\n",
    "def _python_cmd_unzip(source,dest,overwrite=False):\n",
    "    source = Path(source);\n",
    "    dest = Path(dest);\n",
    "    with zipfile.ZipFile(source,'r') as archive:\n",
    "        for member in archive.infolist():\n",
    "            file_path = dest/member.filename\n",
    "            if not file_path.exists():\n",
    "                archive.extract(member, dest)\n",
    "\n",
    "unzipExists = shutil.which(\"unzip\")\n",
    "cmd_unzip = _linux_cmd_unzip if unzipExists else _python_cmd_unzip\n",
    "print(\"using python unzip\" if not unzipExists else \"using cmdline unzip\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y-zFwJIQZSTc"
   },
   "source": [
    "### Gsutil String Handling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gs_str(p:Union[str,Path]):\n",
    "    p = Path(p);\n",
    "    out = \"\"\n",
    "    if is_gcp_path(p):\n",
    "        p = Path(*p.parts[1:])\n",
    "        out = \"gs://\"\n",
    "    out += p.as_posix();\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RtbVZB31p5MW",
    "tags": []
   },
   "source": [
    "### Path Formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SafeDict(dict):\n",
    "    def __missing__(self, key):\n",
    "        return '{' + key + '}'\n",
    "\n",
    "    \n",
    "suffix_regex = ' \\\\${0}'\n",
    "def strsuffix(exp:str,suffixes=None):\n",
    "    if isinstance(suffixes,str):\n",
    "        suffixes = [suffixes]\n",
    "    elif not isinstance(suffixes,Iterable):\n",
    "        suffixes = [\"\\\\S*\"]\n",
    "    for suffix in suffixes:\n",
    "        r = suffix_regex.format(suffix)\n",
    "        exp = re.sub(r,'',exp)\n",
    "    return exp\n",
    "    \n",
    "\n",
    "def format_path(path:Union[str,Path],\n",
    "                collection:Union[bool,None]=None,\n",
    "                movie:Union[int,None]=None,\n",
    "                experiment:Union[str,None]=None,\n",
    "                qc_params:Union[Dict[str,Any],None]=None,\n",
    "                strip_suffix:Union[bool,list[str]]=False,\n",
    "                **kwargs):\n",
    "    is_path = isinstance(path,Path);\n",
    "    path = str(path);\n",
    "    map = SafeDict();\n",
    "    if collection is not None:\n",
    "        map[\"collection\"] = \"using\" if collection else \"no\";\n",
    "    if movie is not None:\n",
    "        map[\"movie\"] = str(movie);\n",
    "    if experiment is not None:\n",
    "        map[\"experiment_nostrip\"] = experiment\n",
    "        if strip_suffix:\n",
    "            experiment = strsuffix(experiment,strip_suffix);\n",
    "        map[\"experiment\"] = experiment;\n",
    "    if qc_params is not None:\n",
    "        map[\"smoothing\"] = [['raw','smoothed'],\n",
    "                             ['spaced','spaced_smoothed']][qc_params.get(\"do_spacing\",False)][qc_params.get(\"do_smoothing\",False)];\n",
    "\n",
    "    map.update(kwargs);\n",
    "    result = path.format_map(map);\n",
    "    if is_path:\n",
    "        result = Path(result);\n",
    "    return result;\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6pukNUHPkGvx",
    "tags": []
   },
   "source": [
    "## Input Paths and Settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OeIEjwglksB5",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Colab Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown where various folders are on disk, you probably never need to change these\n",
    "\n",
    "local_folder = \"content\" #@param {type:\"string\"};\n",
    "local_folder:Path = Path(local_folder);\n",
    "\n",
    "#@markdown local folder where files and folders are stored when downloaded from GCP\n",
    "gcp_transfer_folder = \"gcp_transfer\" #@param {type:\"string\"}\n",
    "gcp_transfer_folder:Path = Path(gcp_transfer_folder)\n",
    "\n",
    "#@markdown local folder where files and folders are stored locally that should be cleared between operations\n",
    "temp_folder = \"temp\" #@param {type:\"string\"}\n",
    "temp_folder:Path = Path(temp_folder)\n",
    "\n",
    "if not os.path.exists(gcp_transfer_folder):\n",
    "  os.mkdir(gcp_transfer_folder);\n",
    "\n",
    "if not os.path.exists(temp_folder):\n",
    "  os.mkdir(temp_folder);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JhBQhuCRgfQs",
    "tags": []
   },
   "source": [
    "### Experimental Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Experiment Name (will be incorporated into output file and folder names)\n",
    "#@markdown Suffixes: if you append an experiment name with \" $[suffix]\", path formatting (with the exception of segmentation analysis) will use the experiment stripped of *suffix\n",
    "experiment = \"2023.3.24 OptoTiam Exp 50\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown Analysis folder: all output analysis data will be output to {analysis output folder}/{experiment name}. Include gs:// to place in the bucket\n",
    "analysis_output_folder = \"gs://optotaxisbucket/Segmentation Analysis\" #@param {type:\"string\"}\n",
    "analysis_output_folder:Path = Path(analysis_output_folder)\n",
    "gcp_analysis_output_folder:Path = analysis_output_folder/experiment if is_gcp_path(analysis_output_folder) else None;\n",
    "local_analysis_output_folder:Path = gcp_transfer_folder/analysis_output_folder.name/experiment;\n",
    "if not os.path.exists(local_analysis_output_folder):\n",
    "    os.makedirs(local_analysis_output_folder);\n",
    "del analysis_output_folder;\n",
    "\n",
    "if \"images_source\" in locals():\n",
    "    prev_im = images_source\n",
    "else:\n",
    "    prev_im = None;\n",
    "\n",
    "if \"cellmasks_source\" in locals():\n",
    "    prev_cm = cellmasks_source\n",
    "else:\n",
    "    prev_cm = None;\n",
    "\n",
    "if \"nucmasks_source\" in locals():\n",
    "    prev_nm = nucmasks_source\n",
    "else:\n",
    "    prev_nm = None;\n",
    "\n",
    "#@markdown For any of the following paths, if you put {experiment} in the string, it'll replace that with the name of the experiment for convenience:\n",
    "\n",
    "#@markdown Location of experiment images. Can be a folder, a .zip file, or a tif stack; If a zip file, should contain a folder of the same name as the zip file. Include 'gs://' to pull from the a google cloud storage bucket\n",
    "images_source_s:str = \"gs://optotaxisbucket/movies/{experiment}/{experiment}\" #@param {type:\"string\"}\n",
    "images_source:Path = Path(format_path(images_source_s,experiment=experiment,strip_suffix=True))\n",
    "\n",
    "#@markdown Location of segmented cell-area masks. Can be a folder, a .zip file, or a tif stack; If a zip file, should contain a folder of the same name as the zip file. Include 'gs://' to pull from the a google cloud storage bucket\n",
    "cellmasks_source_s:str = \"gs://optotaxisbucket/movie_segmentation/{experiment}/segmentation_output_masks/Cell.zip\" #@param {type:\"string\"}\n",
    "cellmasks_source:Path = Path(format_path(cellmasks_source_s,experiment=experiment,strip_suffix=True));\n",
    "\n",
    "#@markdown Location of segmented nucleus-area masks. Can be a folder, a .zip file, or a tif stack; If a zip file, should contain a folder of the same name as the zip file. Include 'gs://' to pull from the a google cloud storage bucket\n",
    "nucmasks_source_s = \"gs://optotaxisbucket/movie_segmentation/{experiment}/segmentation_output_masks/Nucleus.zip\" #@param {type:\"string\"}\n",
    "nucmasks_source:Path = Path(format_path(nucmasks_source_s,experiment=experiment,strip_suffix=True));\n",
    "\n",
    "\n",
    "images_changed,cell_masks_changed,nuc_masks_changed = [False,False,False];\n",
    "if images_source != prev_im:\n",
    "    images_changed = True;\n",
    "if cellmasks_source != prev_cm:\n",
    "    cell_masks_changed = True;\n",
    "if nucmasks_source != prev_nm:\n",
    "    nuc_masks_changed = True;\n",
    "\n",
    "parameters_updated = True\n",
    "\n",
    "\n",
    "for test,name in progress_bar([(images_source,\"Images\"),(cellmasks_source,\"Cell masks\"),(nucmasks_source,\"Nucleus masks\")]):\n",
    "    print(f\"Testing {name}...\")\n",
    "    if is_gcp_path(test):\n",
    "        test = gs_str(test);\n",
    "        if (not str(test).lower().endswith((\".zip\",\".tif\",\".tiff\"))): #directory\n",
    "            test = test + '/*'\n",
    "        valid = os.system(f\"gsutil -q stat \\\"{test}\\\"\");\n",
    "        if (valid != 0): #test location does not exist\n",
    "            raise Exception(f\"Error: {name} dir {test} does not exist in bucket. To ignore this error simply run the succeeding cells.\")\n",
    "    else:\n",
    "        if not os.path.exists(test):\n",
    "            raise Exception(f\"Error: {name} dir {test} does not exist on local machine. To ignore this error simply run the succeeding cells.\")\n",
    "    print(f\"{name} folder verified: {test}\")\n",
    "\n",
    "print(\"Successfully verified inputs - all folders and files exist in the bucket\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SxWIZsq76wHw",
    "tags": []
   },
   "source": [
    "### Analysis Output Filenames\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (2851290609.py, line 12)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  Input \u001b[1;32mIn [1]\u001b[1;36m\u001b[0m\n\u001b[1;33m    metrics_params_path = parameters_folder/\u001b[0m\n\u001b[1;37m                                            ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "#@markdown ## The names/locations of various parameter and analysis output files within the segmentation analysis folder\n",
    "#@markdown ### All input paths are relative to the analysis output folder defined in the previous cell\n",
    "\n",
    "movies_folder = local_analysis_output_folder/Path(\"Output Movies\") #@param\n",
    "\n",
    "#@markdown #### Analysis Parameters\n",
    "parameters_folder = local_analysis_output_folder/\"Analysis Parameters\" #@param\n",
    "if not os.path.exists(parameters_folder):\n",
    "  os.mkdir(parameters_folder)\n",
    "cell_reading_params_path = parameters_folder/\"reading_cells_parameters.pkl\" #@param \n",
    "track_params_path = parameters_folder/\"tracking_parameters.pkl\" #@param\n",
    "metrics_params_path = parameters_folder/\"cellmetrics_parameters.pkl\"\n",
    "qc_params_path = parameters_folder/\"track_qc_output.pkl\" #@param\n",
    "qc_log_path = parameters_folder/\"qc_log.txt\"\n",
    "#make sure no dependency chains\n",
    "del parameters_folder\n",
    "\n",
    "#@markdown Analysis Outputs\n",
    "labeled_cellmasks_path = local_analysis_output_folder/\"labeledmasks.zip\" #@param\n",
    "labeled_nucmasks_path = local_analysis_output_folder/\"labelednucs.zip\" #@param\n",
    "cell_features_path = local_analysis_output_folder/\"cell_features.csv\" #@param\n",
    "raw_tracks_path = local_analysis_output_folder/\"tracks.pkl\" #@param\n",
    "qc_tracks_path = local_analysis_output_folder/\"qc_tracks.pkl\" #@param\n",
    "\n",
    "parameters_updated = False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ej4weULxkO8L",
    "tags": []
   },
   "source": [
    "## Global helper function setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w3WsTzRakQFe",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### File Management"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lHuTdyZu_tRi",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Image Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#keyword should be unique for each type of object pulled from the cloud to avoid collisions\n",
    "def _fetch_image_files(in_path:Path,keyword:str,overwrite:bool,capture:bool)->Path: ##should not be called outside of other helper functions, will always overwrite\n",
    "  if parameters_updated:\n",
    "    raise Exception(\"experimental parameters (file paths, exp name, etc) changed without update analysis outputs. Please run Analysis Output Filenames cell!\")\n",
    "  ctx = contextlib.nullcontext if not capture else capture_output;\n",
    "  is_file = str(in_path).lower().endswith(('.zip','.tif','.tiff'));\n",
    "  is_gcp = is_gcp_path(in_path);\n",
    "  destination = gcp_transfer_folder/keyword/strsuffix(experiment);\n",
    "  if not(os.path.exists(destination)):\n",
    "    os.makedirs(destination);\n",
    "  command_output = None;\n",
    "  if is_gcp and (overwrite or len(os.listdir(destination)) == 0):\n",
    "    with ctx():\n",
    "      if is_file:\n",
    "        command_output = !gsutil -m cp -r \"{gs_str(in_path)}\" \"{gs_str(destination)}\"\n",
    "      else:\n",
    "        command_output = !gsutil -m rsync -r \"{gs_str(in_path)}\" \"{gs_str(destination)}\"\n",
    "    if (command_output is not None and command_output[0].startswith(\"CommandException\")):\n",
    "      raise RuntimeError(f\"Error while downloading {keyword} from bucket: \" + '\\n'.join(command_output));\n",
    "    elif len(os.listdir(destination)) == 0:\n",
    "      raise RuntimeError(\"Error: downloading failed for an unknown reason; downloading command_output:\",command_output);\n",
    "  if (not is_file):\n",
    "    return destination if is_gcp else in_path; #we're done\n",
    "  else:\n",
    "    destination = destination/in_path.name;\n",
    "  if (in_path.suffix == '.zip'):\n",
    "    out_path = destination.with_suffix('');\n",
    "    # print(copy_out);\n",
    "    command_output = None;\n",
    "    if (overwrite or not os.path.exists(out_path)):\n",
    "      with ctx():\n",
    "        command_output = cmd_unzip(destination,destination.parent,overwrite=True)\n",
    "    # print(command_output);\n",
    "    if (command_output is not None and command_output[0].startswith(\"CommandException\")):\n",
    "      raise RuntimeError(f\"Error while unzipping {keyword}: \" + '\\n'.join(command_output));\n",
    "    elif not os.path.exists(out_path):\n",
    "      raise RuntimeError(f\"Error: zip file {destination} does not contain folder {destination.with_suffix('')}\");\n",
    "    return out_path;\n",
    "  elif (in_path.suffix.lower().startswith('.tif')):\n",
    "    raise NotImplementedError(\"unstacking TIF files not yet supported\");\n",
    "  else:\n",
    "    raise NameError(\"Invalid input suffix, input validation should have caught this >:(\");  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fetch_images(force_overwrite=False,capture=False)->Path:\n",
    "  global images_changed;\n",
    "  out = _fetch_image_files(Path(images_source),'images',images_changed or force_overwrite,capture);\n",
    "  images_changed = False;\n",
    "  return out;\n",
    "\n",
    "def fetch_cell_masks(force_overwrite=False,capture=False)->Path:\n",
    "  global cell_masks_changed;\n",
    "  out = _fetch_image_files(Path(cellmasks_source),'cellmasks',cell_masks_changed or force_overwrite,capture);\n",
    "  cell_masks_changed = False;\n",
    "  return out;\n",
    "\n",
    "def fetch_nuc_masks(force_overwrite=False,capture=False)->Path:\n",
    "  global nuc_masks_changed;\n",
    "  out = _fetch_image_files(Path(nucmasks_source),'nucmasks',nuc_masks_changed or force_overwrite,capture);\n",
    "  nuc_masks_changed = False;\n",
    "  return out;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UsFfgymn_wl-"
   },
   "source": [
    "#### Analysis Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def push_analysis(): #VERY IMPORTANT: THIS METHOD RELIES ON local_analysis_output_folder BEING ONE-DEPTH\n",
    "  '''Push to sync the current contents of the gcp bucket's gcp_analysis_output_folder with its local counterpart'''\n",
    "  if parameters_updated:\n",
    "    raise Exception(\"experimental parameters (file paths, exp name, etc) changed without update analysis outputs. Please run Analysis Output Filenames cell!\")\n",
    "  if gcp_analysis_output_folder is None:\n",
    "    print(\"Segmentation analysis is local only; skipping push\");\n",
    "    return;\n",
    "  s = !gsutil -m rsync -r \"{gs_str(local_analysis_output_folder)}\" \"{gs_str(gcp_analysis_output_folder)}\"\n",
    "  if (s[0].startswith(\"CommandException\")):\n",
    "    raise RuntimeError(\"error while uploading analysis folder: \" + '\\n'.join(s));\n",
    "\n",
    "def fetch_analysis():\n",
    "  '''Pull to sync the current contents of local_analysis_output_folder with its counterpart in the gcp bucket'''\n",
    "  if parameters_updated:\n",
    "    raise Exception(\"experimental parameters (file paths, exp name, etc) changed without update analysis outputs. Please run Analysis Output Filenames cell!\")\n",
    "  if gcp_analysis_output_folder is None:\n",
    "    print(\"Segmentation analysis is local only; skipping fetch\");\n",
    "    return;\n",
    "  if not os.path.exists(local_analysis_output_folder):\n",
    "    os.makedirs(local_analysis_output_folder);\n",
    "  t = !gsutil ls \"{gs_str(gcp_analysis_output_folder)}\"\n",
    "  if (t[0].startswith(\"CommandException\")):\n",
    "    print(\"analysis folder not found in bucket, skipping analysis fetch:\");\n",
    "    print('\\n'.join(t));\n",
    "    return;\n",
    "  s = !gsutil -m rsync -r \"{gs_str(gcp_analysis_output_folder)}\" \"{gs_str(local_analysis_output_folder)}\"\n",
    "  if (s[0].startswith(\"CommandException\")):\n",
    "    raise RuntimeError(\"error while downloading analysis folder: \" + '\\n'.join(s));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _submit_masks(source:Path,dest:Path):\n",
    "  cmd_zip(source,dest,recurse=True);\n",
    "  push_analysis()\n",
    "\n",
    "def submit_labeled_cellmasks(submission:Path):\n",
    "  '''Input a folder containing a list of cell masks; will zip, put into analysis, and push'''\n",
    "  return _submit_masks(submission,labeled_cellmasks_path)\n",
    "\n",
    "def submit_labeled_nucmasks(submission:Path):\n",
    "  '''Input a folder containing a list of nucleus masks; will zip, put into analysis, and push'''\n",
    "  return _submit_masks(submission,labeled_nucmasks_path);\n",
    "\n",
    "def _fetch_masks(file:Path): #THIS NEEDS TO BE TESTED!!!!!!!!!\n",
    "  fetch_analysis()\n",
    "  dest = temp_folder/(file.stem)\n",
    "  cmd_unzip(file,dest,overwrite=True);\n",
    "  return dest\n",
    "\n",
    "def fetch_labeled_cellmasks():\n",
    "  '''Pulls and unzips labeled and filtered cell masks from segmentation analysis; returns a folder containing the masks'''\n",
    "  return _fetch_masks(labeled_cellmasks_path);\n",
    "\n",
    "def fetch_labeled_nucmasks():\n",
    "  '''Pulls and unzips labeled and filtered nucleus masks from segmentation analysis; returns a folder containing the masks'''\n",
    "  return _fetch_masks(labeled_nucmasks_path);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zBs3iPXU_0KH"
   },
   "source": [
    "#### Misc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_rm_error( func, path, exc_info):\n",
    "    # path contains the path of the file that couldn't be removed\n",
    "    # let's just assume that it's read-only and unlink it.\n",
    "    os.chmod( path, stat.S_IWRITE )\n",
    "    # os.unlink( path )\n",
    "\n",
    "def cleardir(dir): #clears all files in dir without deleting dir\n",
    "  for f in os.scandir(dir):\n",
    "    # f = os.path.join(dir,f)\n",
    "    if os.path.isdir(f): shutil.rmtree(f,onerror=on_rm_error); #just in case\n",
    "    else: os.remove(f);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4FjItCTOiau9"
   },
   "source": [
    "### Cell filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getcells(filecell:Union[Union[str, bytes, os.PathLike],np.ndarray],filenuc:Union[Union[str, bytes, os.PathLike],np.ndarray],parameters,return_metrics=False,centertype=\"approximate-medoird\"):\n",
    "  #membrane\n",
    "  maskMem:np.ndarray=imread(filecell) if not isinstance(filecell,np.ndarray) else filecell;\n",
    "  # print(np.min(maskmem));\n",
    "  # print(np.max(maskmem));\n",
    "  maskMem[maskMem>0]=1\n",
    "  #fill holes\n",
    "  maskMem=ndimage.binary_fill_holes(maskMem).astype(int);\n",
    "\n",
    "  #nuclei\n",
    "  maskNuc:np.ndarray = imread(filenuc) if not isinstance(filenuc,np.ndarray) else filenuc;\n",
    "  maskNuc[maskNuc>0]=1\n",
    "\n",
    "  #label different objectes in masks\n",
    "  maskMem,numMem = measure.label(maskMem,return_num=True)\n",
    "  maskNuc,numNuc = measure.label(maskNuc,return_num=True);\n",
    "\n",
    "  if numMem < 255 and numNuc < 255:\n",
    "    maskMem = maskMem.astype('uint8');\n",
    "    maskNuc = maskNuc.astype('uint8');\n",
    "  else:\n",
    "    maskMem = maskMem.astype('uint16');\n",
    "    maskNuc = maskNuc.astype('uint16');\n",
    "  \n",
    "  #FILTERS\n",
    "  if parameters['remove_cells_touching_edge'] == True:\n",
    "    maskMem=remove_touching_edge(maskMem)\n",
    "  \n",
    "  if parameters['filter_cell_size'] == True:\n",
    "    maskMem = morphology.remove_small_objects(maskMem, parameters['minareacell'])       \n",
    "    maskMem = remove_large_objects(maskMem, parameters['maxareacell'])\n",
    "  \n",
    "  if parameters['filter_nuc_size'] == True:\n",
    "    maskNuc = morphology.remove_small_objects(maskNuc, parameters['minareanuc'])       \n",
    "  \n",
    "  if parameters['remove_multi_nuclei_cells'] == True:\n",
    "    maskMem = remove_multiple_nuclei_cells(maskMem,maskNuc)\n",
    "\n",
    "  if (return_metrics):\n",
    "    #if there are cells get metrics\n",
    "    ids=list(range(1,numMem+1));\n",
    "    #remove 0 (background) from ids\n",
    "    # ids.remove(0)\n",
    "    if len(ids) > 0:\n",
    "      cellsmetrics = measure.regionprops_table(maskMem, properties=('label','area'))\n",
    "      cellsmetrics=pd.DataFrame(cellsmetrics)\n",
    "      if (len(cellsmetrics['label']) > 0 and len(cellsmetrics['area']) > 0):\n",
    "\n",
    "        #GET CENTERS\n",
    "        #get labels\n",
    "        labels=cellsmetrics['label']    \n",
    "        #Because 'label' was copied from the table, after computing the centers \n",
    "        #and concatenating them to the table they should be in the right order\n",
    "        if not all_centers:\n",
    "            centers=get_centers(maskMem,centertype,labels, False)\n",
    "            #add centers to cell properties\n",
    "            cntr=pd.DataFrame(data=np.asarray(centers),columns=[centertype+'x',centertype+'y'])\n",
    "            cellsmetrics=pd.concat([cellsmetrics,cntr],axis=1)\n",
    "        else:\n",
    "            for ctype in get_centers.valid_centers: \n",
    "                ## to save on import complexity, I (H) just set get_centers.valid_centers to the list of valid center inputs. \n",
    "                ## It's still a function, it just has the attribute now. source is libraries.centers\n",
    "                centers=get_centers(maskMem,ctype,labels,False)\n",
    "                #add centers to cell properties\n",
    "                cntr=pd.DataFrame(data=np.asarray(centers),columns=[centertype+'x',centertype+'y'])\n",
    "                cellsmetrics=pd.concat([cellsmetrics,cntr],axis=1)\n",
    "        \n",
    "        # centers=get_centers(maskMem,'centroid',labels, False)\n",
    "        # #add centers to cell properties\n",
    "        # centroid=pd.DataFrame(data=np.asarray(centers),columns=['centroidx','centroidy'])\n",
    "        # cellsmetrics=pd.concat([cellsmetrics,centroid],axis=1)\n",
    "      \n",
    "        \n",
    "\n",
    "\n",
    "    else:\n",
    "      cellsmetrics=pd.DataFrame();\n",
    "    return cellsmetrics, maskMem, maskNuc \n",
    "  else:\n",
    "    return maskMem,maskNuc\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f-PoRJRvn9zd",
    "tags": []
   },
   "source": [
    "# Input Movie Reading Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8skHjf8JoCmn",
    "tags": []
   },
   "source": [
    "### Copy images and masks from GCP to Colab (May take some time, zip folders in GCP to speed this up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "folderimages = fetch_images();\n",
    "foldermasks = fetch_cell_masks();\n",
    "try:\n",
    "    foldernucmasks = fetch_nuc_masks();\n",
    "except Exception as e:\n",
    "    foldernucmasks = None;\n",
    "    raise Exception(\"Warning: nucleus masks not found. If you wish to visualize anyway, keep running cells in Segmentation quality control, but cell filtering will be disabled and settings will not save\") from e;\n",
    "clear_output();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fetch_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CrQWf9_GmVTF"
   },
   "source": [
    "### Initialize parameters for reading images and upload (select movies and frame ranges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yqTMSYzCZSUX",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Identify movies and frames in input image folder (code hidden)\n",
    "imagenames=[f for f in os.listdir(foldermasks) if f.endswith('.TIF')]\n",
    "#Get list of movies\n",
    "movies = [int(re.findall(r\"s(\\d+).\", imagenames[i])[0]) for i in range(len(imagenames)) ]\n",
    "#get set of unique elements\n",
    "movies = list(set(movies))\n",
    "#sort\n",
    "movies.sort()\n",
    "print(\"movies detected in experiment:\",movies);\n",
    "#Get frame numbers and show the largest number\n",
    "frames = {m:[] for m in movies};\n",
    "for name in imagenames:\n",
    "  match = re.findall(r\"s(\\d+)_t(\\d+).\", name)[0];\n",
    "  # print(match);\n",
    "  frames[int(match[0])].append(int(match[1]));\n",
    "for m in movies:\n",
    "  frames[m].sort();\n",
    "# nframes={m:max(f) for m,f in frames.items()}\n",
    "# print(\"frames detected per movie:\",nframes);\n",
    "# print(frames);\n",
    "# imagenames[0]\n",
    "basename= re.findall(r\"(.*)_s\",imagenames[0])[0]\n",
    "print(\"image basename:\",basename);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CWFqLn4vZSUc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Set movie frame ranges\n",
    "use_frame_range = False #@param {\"type\":\"boolean\"}\n",
    "\n",
    "#@markdown Enter frame ranges: dictionary of movie:[rangemin,rangemax] items in raw python syntax\n",
    "frame_ranges:Dict =  {8:(0,145)}#@param {\"type\":\"raw\"}\n",
    "\n",
    "#@markdown whether to exclude specific movies from the list of movies\n",
    "do_exclude_movies = False #@param {\"type\":\"boolean\"}\n",
    "\n",
    "from typing import Iterable\n",
    "#@markdown list of movies to exclude\n",
    "exclude_movies:Iterable = [] #@param {\"type\":\"raw\"}\n",
    "\n",
    "if do_exclude_movies:\n",
    "  [movies.remove(m) for m in exclude_movies]\n",
    "  print(\"updated movies:\",movies)\n",
    "\n",
    "if use_frame_range:\n",
    "  new_frames = {m:frames[m] if m not in frame_ranges else list(range(max(min(frames[m]),frame_ranges[m][0]),1+min(max(frames[m]),frame_ranges[m][1]))) for m in movies};\n",
    "  assert set(new_frames.keys()) == set(movies);\n",
    "  frames = new_frames;\n",
    "  print(\"updated frame ranges:\",frames)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fLsXzIQ9uS_4"
   },
   "source": [
    "### Input and Save Cell filtering parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "21NB9KD8ZSUi",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Input and save cell filtering parameters\n",
    "\n",
    "\n",
    "#@markdown #### Cell location parameters\n",
    "\n",
    "remove_cells_touching_edge = True #@param {type:\"boolean\"}\n",
    "\n",
    "remove_multi_nuclei_cells = False #@param {type:\"boolean\"}\n",
    "\n",
    "\n",
    "#@markdown #### Set size bounds (number of pixels)\n",
    "\n",
    "filter_cell_size = True #@param {type:\"boolean\"}\n",
    "\n",
    "\n",
    "#minarea=600\n",
    "#maxarea=3200\n",
    "#nucminarea=21*pixelsize**2\n",
    "\n",
    "minareacell=650 #@param {type:\"integer\"}\n",
    "maxareacell=25974 #@param {type:\"integer\"}\n",
    "\n",
    "filter_nuc_size = True #@param {type:\"boolean\"}\n",
    "minareanuc=5 #@param {type:\"integer\"}\n",
    "\n",
    "reading_cells_parameters = {\n",
    "    'experiment':experiment, \n",
    "    'basename':basename,\n",
    "    'movies':movies,\n",
    "    'frames':frames,\n",
    "    'remove_cells_touching_edge':remove_cells_touching_edge,\n",
    "    'remove_multi_nuclei_cells':remove_multi_nuclei_cells,\n",
    "    'filter_cell_size':filter_cell_size,\n",
    "    'minareacell':minareacell,\n",
    "    'maxareacell':maxareacell,\n",
    "    'filter_nuc_size':filter_nuc_size,\n",
    "    'minareanuc':minareanuc,\n",
    "    # 'img_min_intensity':minintensity,\n",
    "    # 'img_max_intensity':maxintensity\n",
    "              }\n",
    "\n",
    "if not os.path.exists(local_analysis_output_folder):\n",
    "  os.makedirs(local_analysis_output_folder)\n",
    "\n",
    "with open(cell_reading_params_path, 'wb') as handle:\n",
    "    pickle.dump(reading_cells_parameters, handle)\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4JsH0gJkeNO-",
    "tags": []
   },
   "source": [
    "# [Optional] Segmentation Visualization - Segmentation quality control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J8Fh0Q29fWSf"
   },
   "source": [
    "### Copy images and masks from GCP to Colab (May take some time, zip folders in GCP to speed this up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I7kUnzlrZSUp"
   },
   "outputs": [],
   "source": [
    "folderimages = fetch_images();\n",
    "foldermasks = fetch_cell_masks();\n",
    "try:\n",
    "  foldernucmasks = fetch_nuc_masks();\n",
    "except Exception as e:\n",
    "  foldernucmasks = None;\n",
    "  raise Exception(\"Warning: nucleus masks not found. If you wish to visualize anyway, keep running cells in Segmentation quality control, but cell filtering will be disabled and settings will not save\") from e;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fVsxX0fvZSUq"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "with open(cell_reading_params_path,'rb') as f:\n",
    "  reading_cells_parameters = pickle.load(f);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CxBtSBqAEqoX"
   },
   "source": [
    "### Fetch Analysis, Load Movie Reading Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "phlpgrX5ZSUu"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "\n",
    "with open(cell_reading_params_path,'rb') as f:\n",
    "  reading_cells_parameters = pickle.load(f);\n",
    "\n",
    "basename = reading_cells_parameters['basename'];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lg_G4zagKEvF"
   },
   "source": [
    "## Load separate cell masks and nucleus masks and visualize together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G-V64-3iZSU0"
   },
   "outputs": [],
   "source": [
    "#@markdown ##Compute Movie Frames\n",
    "\n",
    "#@markdown Which movie to view\n",
    "movie=4 #@param {type:\"integer\"}\n",
    "if movie not in reading_cells_parameters['movies']:\n",
    "  raise Exception(f\"invalid movie number, check that it's in range hasn't been excluded from the experiment.\\nMovie number {movie} is not in experiment with movies: {movies}\")\n",
    "\n",
    "#@markdown Frame downsample rate (how often a frame is shown); increase this number for a faster processing time but less fine time-resolution\n",
    "frspace=3 #@param {type:\"integer\"}\n",
    "\n",
    "#@markdown Nucleus intensity for combined image (between 0 and 1)\n",
    "mfrac=0.7 #@param {type:\"number\"}\n",
    "\n",
    "\n",
    "#@markdown how much to resize the image by (the more downscaled, the faster the playing but the less the resolution)\n",
    "downfrac = 0.5 #@param {type:\"number\"};\n",
    "\n",
    "\n",
    "\n",
    "#@markdown whether to show the splitting sections on the image\n",
    "show_splitting_sections = False #@param{type:\"boolean\"}\n",
    "\n",
    "splitting_overlay_opacity = 0.3;\n",
    "splitting_overlay_contrast = 1;\n",
    "\n",
    "#@markdown ### If showing splitting sections, input the splitting parameters used\n",
    "#@markdown x_slices, y_slices are how many pieces in each direction the image is split into\n",
    "x_slices = 5 #@param {type:\"integer\"}\n",
    "y_slices =  5 #@param {type:\"integer\"}\n",
    "#@markdown dx, dy are the extra context around the segmented center in both directions\n",
    "dx =  42#@param {type:\"integer\"}\n",
    "dy =  32#@param {type:\"integer\"}\n",
    "#@markdown x and y crop are how much to straight remove from the image to make the sizes able to be subdivided nicely\n",
    "x_crop = 0 #@param {type:\"integer\"}\n",
    "y_crop = 0 #@param {type:\"integer\"}\n",
    "\n",
    "context_bounds = [dy,dx]*2 #assuming x and y symmetrical, not always true -- fix?\n",
    "crop = [y_crop,x_crop]*2 #both of these are negative y, negative x, positive y, positive x\n",
    "splitting_overlay = None;\n",
    "if show_splitting_sections:\n",
    "  sh = imread(folderimages/[f for f in os.listdir(folderimages) if f.endswith(\".TIF\")][0] or f.endswith('.tif')).shape;\n",
    "  im = np.ndarray((sh[0],sh[1]));\n",
    "  im.fill(0);\n",
    "  print(im.shape);\n",
    "  M = (im.shape[0]-context_bounds[0]-context_bounds[2]-crop[0]-crop[2])/y_slices;\n",
    "  N = (im.shape[1]-context_bounds[1]-context_bounds[3]-crop[1]-crop[3])/x_slices;\n",
    "  print(M,N);\n",
    "  if int(M) != M or int(N) != N:\n",
    "    raise Exception(f\"ERROR: invalid image splitting parameters for image of shape {im.shape}\");\n",
    "  M = int(M)\n",
    "  N = int(N)\n",
    "  tiles = [im[y-context_bounds[0]:y+M+context_bounds[2],x-context_bounds[1]:x+N+context_bounds[3]] \n",
    "    for y in range(context_bounds[0]+crop[0],im.shape[0]-crop[0]-crop[2]-context_bounds[0]-context_bounds[2],M) \n",
    "    for x in range(context_bounds[1]+crop[1],im.shape[1]-crop[1]-crop[3]-context_bounds[1]-context_bounds[3],N)];\n",
    "\n",
    "  outtiles = [];\n",
    "  for i,c in enumerate(tiles):\n",
    "    tile = c.copy();\n",
    "    if (i % 2 == 0):\n",
    "      tile.fill(splitting_overlay_contrast);\n",
    "    outtiles.append(tile);\n",
    "  tiles = outtiles;\n",
    "\n",
    "  stitchMasks = [];\n",
    "  for i,m in enumerate(tiles):\n",
    "    x = i // y_slices;\n",
    "    y = i % x_slices;\n",
    "    imBounds = [crop[0]+context_bounds[0] if x != 0 else 0,m.shape[0]-crop[2]-context_bounds[2] if x != x_slices-1 else m.shape[0],crop[1]+context_bounds[1] if y!= 0 else 0 ,m.shape[1]-crop[3]-context_bounds[3] if y != y_slices - 1 else m.shape[1]];\n",
    "    stitchMasks.append(m[imBounds[0]:imBounds[1],imBounds[2]:imBounds[3]]);\n",
    "  splitting_overlay = np.concatenate([np.concatenate(stitchMasks[i*x_slices:(i+1)*x_slices],axis=1) for i in range(y_slices)]);\n",
    "  print(splitting_overlay.shape);\n",
    "\n",
    "frames = reading_cells_parameters['frames'];\n",
    "spacedframes = frames[movie][::frspace] #don't look at every frame for QC\n",
    "# print(frames);\n",
    "combineds=[]\n",
    "for i in progress_bar(spacedframes):\n",
    "    filename_img = basename + \"_s\" + str(movie)+\"_t\"+str(i)+'.tif';\n",
    "    filename_mask = basename + \"_s\" + str(movie)+\"_t\"+str(i)+'.TIF';\n",
    "\n",
    "    cell_path = foldermasks/filename_mask;\n",
    "    nuc_path = foldernucmasks/filename_mask;\n",
    "\n",
    "    maskmem,masknuc = getcells(cell_path,nuc_path,reading_cells_parameters,return_metrics=False);\n",
    "    \n",
    "    #read image\n",
    "    imageorig = imread(folderimages/filename_img)\n",
    "    #rescale image intensity\n",
    "    image = rescale_intensity(imageorig);\n",
    "    #rescale image\n",
    "    image=resize(image, (maskmem.shape[0] * downfrac , maskmem.shape[1] * downfrac),  anti_aliasing=True)\n",
    "    \n",
    "    #unlabel for visualization\n",
    "    maskmem[maskmem>0]=1\n",
    "    masknuc[masknuc>0]=1\n",
    "    maskmem = rescale_intensity(maskmem.astype(image.dtype));\n",
    "    masknuc = rescale_intensity(masknuc.astype(image.dtype));\n",
    "\n",
    "    if show_splitting_sections:\n",
    "      maskcomb = (maskmem - mfrac*masknuc + splitting_overlay_opacity*splitting_overlay);\n",
    "    else:\n",
    "      maskcomb = (maskmem - mfrac*masknuc)\n",
    "    maskcomb[maskcomb<0]=0; #this is so stupid\n",
    "\n",
    "    maskcomb=resize(maskcomb,(int(maskcomb.shape[0]*downfrac), int(maskcomb.shape[1]*downfrac)),  anti_aliasing=True )\n",
    "\n",
    "    combined = np.hstack((image,maskcomb))\n",
    "    combined = rescale_intensity(combined,out_range=np.uint16).astype(np.uint16);\n",
    "\n",
    "    combineds.append(combined)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A4CIfmkwZSU-"
   },
   "outputs": [],
   "source": [
    "#@markdown ## Play movie using mediapy\n",
    "\n",
    "#@markdown Output Video Framerate (fps)\n",
    "framerate =  2#@param {type:\"number\"}\n",
    "\n",
    "media.show_video(progress_bar(combineds),fps=framerate,qp=5);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "to3eeDC9ZSVB"
   },
   "outputs": [],
   "source": [
    "#@markdown ## Scroll through movie using matplotlib\n",
    "with capture_output():\n",
    "  def f(frame):\n",
    "      print(\"displaying\",frame);\n",
    "      l.set_data(combineds[frame-1])\n",
    "      fig.canvas.draw() #use with %matplotlib notebook\n",
    "      display(fig) #use with %matplotlib inline\n",
    "      # plt.imshow(img_median, cmap=\"gray\") #This would work\n",
    "      #plt.show()\n",
    "\n",
    "  fig = plt.figure(figsize=(14, 8))\n",
    "  ax_mask = fig.add_subplot(111) \n",
    "  l = ax_mask.imshow(combineds[0], cmap=\"gray\")\n",
    "\n",
    "interact(f, frame=widgets.IntSlider(min=1,max=len(combineds),step=1,value=1,msg_throttle=1,layout=Layout(width='90%', height='40px'))) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xe9EIU8yd3Lt"
   },
   "source": [
    "### Save Movie as .tiff stack (full quality)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GOSQhqOXZSVF"
   },
   "outputs": [],
   "source": [
    "#@markdown Where to save the output movie file, locally (relative/absolute) or in the GCP bucket (include the gs:// for bucket path);\n",
    "movie_path =  f\"optotaxisbucket/QC_movies/{experiment}-s{movie}.tiff\" #@param \n",
    "movie_path = Path(movie_path);\n",
    "# if not str(gcp_movie_path).lower().endswith((\".tif\",\".tiff\")):\n",
    "#   gcp_movie_path = gcp_movie_path.with_suffix(\".tiff\");\n",
    "cleardir(temp_folder);\n",
    "gcp_movie = is_gcp_path(movie_path)\n",
    "local_output = movie_path\n",
    "if gcp_movie:\n",
    "    local_output = temp_folder/movie_path.name;\n",
    "\n",
    "\n",
    "with tifffile.TiffWriter(local_output) as stack:\n",
    "    for im in combineds: #case \n",
    "        stack.save(\n",
    "            im,\n",
    "            photometric='minisblack', \n",
    "            contiguous=True\n",
    "        )\n",
    "\n",
    "if gcp_movie:\n",
    "    !gsutil -m cp \"{gs_str(local_output)}\" \"{gs_str(gcp_movie_path)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yX_KR34g5iPt"
   },
   "source": [
    "### Save Movie as .mp4 (compressed) video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zmAKJJ7hZSVI"
   },
   "outputs": [],
   "source": [
    "#@markdown Where to save the output movie file in the GCP bucket (include the gs://);\n",
    "gcp_movie_path =  f\"optotaxisbucket/QC_movies/splitting_lines_{experiment}-s{movie}.mp4\" #@param \n",
    "gcp_movie_path = Path(gcp_movie_path);\n",
    "\n",
    "#@markdown Output Video Framerate\n",
    "framerate = 2 #@param {type:\"number\"}\n",
    "# if not str(gcp_movie_path).lower().endswith((\".tif\",\".tiff\")):\n",
    "#   gcp_movie_path = gcp_movie_path.with_suffix(\".tiff\");\n",
    "cleardir(temp_folder);\n",
    "local_output = temp_folder/gcp_movie_path.name;\n",
    "print(\"writing video file...\");\n",
    "media.write_video(local_output,combineds,fps=framerate);\n",
    "\n",
    "!gsutil -m cp \"{gs_str(local_output)}\" \"{gs_str(gcp_movie_path)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iESuwsddgch5",
    "tags": []
   },
   "source": [
    "# Get cell features, labeled masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4ph4g4faowqF",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Clear temporary directory for use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OkNiNBz_ZSVN"
   },
   "outputs": [],
   "source": [
    "cleardir(temp_folder);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pLQopD0KU1Kw",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Fetch cell and nucleus masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0EbswVqeZSVP"
   },
   "outputs": [],
   "source": [
    "folder_incell = fetch_cell_masks();\n",
    "folder_innuc = fetch_nuc_masks();\n",
    "\n",
    "#get cell_nuc_masks names\n",
    "masknames=[f for f in os.listdir(folder_incell) if f.endswith('.TIF')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0I1eviuZMGrW"
   },
   "source": [
    "## Create Cell Metrics (centers and labels), Save data in analysis_output_folder, and transfer to GCP bucket. (~1-2 minutes per movie)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KoffbfqiZSVT"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "if not os.path.exists(local_analysis_output_folder):\n",
    "  os.makedirs(local_analysis_output_folder)\n",
    "\n",
    "centertype ='approximate-medoid' #@param [\"approximate-medoid\",\"centroid\",\"medoid\",\"ellipse\",\"largest-circle\"]\n",
    "metrics_params = {\"centroidtype\":centertype}\n",
    "with open(metrics_params_path,'wb') as handle:\n",
    "    pickle.dump(metrics_params,handle)\n",
    "\n",
    "\n",
    "#SAVE LABELED MASK SO THAT FEATURES CAN BE EXTRACTED AFTER TRACKING\n",
    "folder_cellmasks_labeled = temp_folder/labeled_cellmasks_path.stem\n",
    "folder_nucmasks_labeled = temp_folder/labeled_nucmasks_path.stem\n",
    "if not os.path.exists(folder_cellmasks_labeled):\n",
    "  os.mkdir(folder_cellmasks_labeled);\n",
    "if not os.path.exists(folder_nucmasks_labeled):\n",
    "  os.mkdir(folder_nucmasks_labeled);\n",
    "try:\n",
    "  with open(cell_reading_params_path,'rb') as f:\n",
    "    reading_cells_parameters = pickle.load(f);\n",
    "except Exception as e:\n",
    "  raise Exception(\"Error: missing reading cells parameters, did you make sure to run segmentation quality control first?\") from e\n",
    "\n",
    "basename = reading_cells_parameters['basename'];\n",
    "frames = reading_cells_parameters['frames'];\n",
    "movies = reading_cells_parameters['movies'];\n",
    "\n",
    "#not necessary to specify columns, but this way 'movie' and 'frame' will be at the beginning of the cells data frame\n",
    "cellmetrics=pd.DataFrame(columns=['movie','frame'],dtype=int);\n",
    "metriclist = [];\n",
    "m = master_bar(movies)\n",
    "for s in m:\n",
    "  # print(s,frames[s]);\n",
    "  for fr in progress_bar(frames[s],parent=m):\n",
    "      filename = folder_incell/(basename + \"_s\" + str(s)+'_t'+str(fr)+'.TIF');\n",
    "      filenamenuc = folder_innuc/(basename + \"_s\" + str(s)+'_t'+str(fr)+'.TIF');\n",
    "      \n",
    "      mem = imread(filename);\n",
    "      mem = ndimage.binary_fill_holes(mem);\n",
    "      mem[mem != 0] = 1;\n",
    "      nuc = imread(filenamenuc);\n",
    "      nuc[nuc != 0] = 1;\n",
    "      \n",
    "      #get cell features in a mask and the labeled mask\n",
    "      cellmetric, labeledcellmask,labelednucmask = getcells(mem,nuc,reading_cells_parameters,return_metrics=True,centertype=centertype);\n",
    "\n",
    "\n",
    "      #SAVE LABELED MASK SO THAT FEATURES CAN BE EXTRACTED AFTER TRACKING\n",
    "      imsave(folder_cellmasks_labeled/(basename + \"_s\" + str(s)+'_t'+str(fr)+'.TIF') , labeledcellmask,check_contrast=False);\n",
    "      imsave(folder_nucmasks_labeled/(basename + \"_s\" + str(s)+'_t'+str(fr)+'.TIF') , labelednucmask,check_contrast=False);\n",
    "      if not cellmetric.empty:\n",
    "        # print(\"hello\");\n",
    "        #Add columns corresponding to movie and frame\n",
    "        cellmetric['movie']=s\n",
    "        cellmetric['frame']=fr\n",
    "        #adds metricsmaskcells below cells\n",
    "        metriclist.append(cellmetric)\n",
    "\n",
    "cellmetrics = pd.concat(metriclist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JTsI20syZSVX"
   },
   "outputs": [],
   "source": [
    "cellmetrics.to_csv(local_analysis_output_folder/'cell_features.csv')\n",
    "submit_labeled_cellmasks(folder_cellmasks_labeled);\n",
    "submit_labeled_nucmasks(folder_nucmasks_labeled)\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zdunkzbKZSVZ"
   },
   "outputs": [],
   "source": [
    "cleardir(temp_folder);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dp8hf3YmBiqC",
    "tags": []
   },
   "source": [
    "# Track Cells (Requires Cell Features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z_X2D_NUmzxV"
   },
   "source": [
    "Reads cell_features from folder_output in GCP bucket, writes tracking_parameters and tracks to folder_output in GCP bucket\n",
    "\n",
    " can be either \"centroid\" or \"approximate-medoid\"\n",
    "\n",
    "maxdist is the maximum distance in pixels allowed for a cell jump\n",
    "\n",
    "dfracsize is the maximum fractional change in area allowed\n",
    "\n",
    "The output (tracks) should be a file in the output folder ending with \"tracks.pkl\" containing a list of tables (each table is a track)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lr3qIC-AZSVg"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "\n",
    "with open(cell_reading_params_path,'rb') as f:\n",
    "    reading_cells_parameters = pickle.load(f);\n",
    "if os.path.exists(metrics_params_path):\n",
    "    with open(metrics_params_path,'rb') as f:\n",
    "        metrics_params = pickle.load(f)\n",
    "else:\n",
    "    if os.path.exists(track_params_path):\n",
    "        with open(track_params_path,'rb') as f:\n",
    "            metrics_params = {\"centroidtype\":pickle.load(f)[\"centroidype\"]}\n",
    "    else:\n",
    "        metrics_params = {\"centroidtype\":\"approximate-medoid\"}\n",
    "\n",
    "#@markdown # Cell Tracking Parameters\n",
    "\n",
    "#@markdown which center type (centroid or medoid) to use for tracking\n",
    "centroidtype = metrics_params[\"centroidtype\"]\n",
    "#@markdown maximum jump size allowed in pixels\n",
    "max_tracked_speed =  300#@param {type:\"integer\"}\n",
    "#@markdown maximum fractional area change \n",
    "dfracsize=0.9 #@param {type:\"number\"}\n",
    "\n",
    "\n",
    "\n",
    "#@markdown ## Disappeared Cell Tracking Parameters\n",
    "#@markdown #### Filters on how a track whose that has disappeared can snap to a newly appeared cell\n",
    "\n",
    "#@markdown maximum frames the tracker will remember a cell's position when not in frame before it gets deleted\n",
    "max_track_disappeared_time =  10 #@param {type:\"integer\"};\n",
    "\n",
    "#@markdown minimum length of continuous track required for persistence through disappeared\n",
    "min_track_persistence_time = 5 #@param {type:\"integer\"}\n",
    "\n",
    "options = [\"Speed (linear)\", \"Diffusivity (MSD)\"]\n",
    "class SpeedType(Enum): \n",
    "  linear = \"linear\";\n",
    "  MSD = \"MSD\";\n",
    "  def __str__(self): \n",
    "    return self.value;\n",
    "\n",
    "#@markdown Speed type - whether to use linear speed or mean squared displacement to filter for the distance a cell can move while disappeared\n",
    "untracked_speed_type_str = \"Diffusivity (MSD)\" #@param [\"Speed (linear)\", \"Diffusivity (MSD)\"]\n",
    "untracked_speed_type = [SpeedType.linear,SpeedType.MSD][options.index(untracked_speed_type_str)];\n",
    "\n",
    "#@markdown If using linear speed: maximum average distance a cell that has disappeared can move per frame while still being snapped back to nearest previous position; that is, cell is accepted if distance(C1,C2) < time*max_untracked_speed\n",
    "max_untracked_speed = 500 #@param {type:\"number\"}\n",
    "\n",
    "#@markdown If using MSD: maximum MSD of an cell that has disappeared moving over the course of t frames to be considered the same as a previously tracked cell; that is, cells will be considered the same if distance(C1,C2)^2 < time*max_untracked_diffusivity. First frame's distance will be sqrt(D), then sqrt(2D), then sqrt(3D), etc\n",
    "max_untracked_diffusivity = 36000 #@param {type:\"number\"}\n",
    "\n",
    "\n",
    "#@markdown ### Cell Collection: Whether to have only a specific section of movie where new tracks can be added\n",
    "#@markdown Probably outdated, initial track delay in track qc is probably better, which does the same thing but throughout the movie with after-the-fact information\n",
    "do_cell_collection = False #@param {type:\"boolean\"}\n",
    "#@markdown Cell Collection Time: For how many frames at the begininng of the movie do we still accept new cells\n",
    "cell_collection_time = 36 #@param {type:\"integer\"}\n",
    "\n",
    "\n",
    "trackinging_parameters={\n",
    "    'centroidtype':centroidtype, \n",
    "    'max_tracked_speed':max_tracked_speed, \n",
    "    'dfracsize':dfracsize,\n",
    "    'max_disappeared_time':max_track_disappeared_time,\n",
    "    'untracked_speed_type':str(untracked_speed_type),\n",
    "    'do_cell_collection':do_cell_collection,\n",
    "    'cell_collection_time':cell_collection_time,\n",
    "    'min_track_persistence_time':min_track_persistence_time};\n",
    "if untracked_speed_type == SpeedType.linear:\n",
    "  trackinging_parameters['max_untracked_speed'] = max_untracked_speed;\n",
    "elif untracked_speed_type == SpeedType.MSD:\n",
    "  trackinging_parameters['max_untracked_diffusivity'] = max_untracked_diffusivity;\n",
    "\n",
    "\n",
    "with open(track_params_path, 'wb') as handle:\n",
    "    pickle.dump(trackinging_parameters, handle)\n",
    "\n",
    "#This specify the cell metrics tracking is going to use\n",
    "fields = ['label','movie','frame',centroidtype+'x',centroidtype+'y','area']\n",
    "\n",
    "#copy cell shape features from bucket to colab\n",
    "push_analysis();\n",
    "\n",
    "#read cell shape features as data frame\n",
    "sample_cells:pd.DataFrame = pd.read_csv(local_analysis_output_folder/'cell_features.csv', usecols=fields)\n",
    "\n",
    "sample:Dict[int,Dict[int,pd.DataFrame]] = {}; #movie index, tracks\n",
    "# to_use = [m for m in movies if m not in fails] if 'fails' in locals() else movies;\n",
    "\n",
    "class Cell:\n",
    "  def __init__(self,frame:pd.DataFrame):\n",
    "    self.frame = frame;\n",
    "    self.min_dist = None;\n",
    "  \n",
    "  @staticmethod\n",
    "  def cell_distance(t0:Cell,t1:Cell)->float:\n",
    "      t0_f = t0.frame.iloc[0];\n",
    "      t1_f = t1.frame.iloc[0];\n",
    "      out = math.sqrt((t0_f[centroidtype+'x']-t1_f[centroidtype+'x'])**2 + (t0_f[centroidtype+'y']-t1_f[centroidtype+'y'])**2);\n",
    "      if (t1.min_dist is None or out < t1.min_dist):\n",
    "        t1.min_dist = out;\n",
    "      return out;\n",
    "\n",
    "  def __str__(self):\n",
    "    return f\"Cell <{builtins.id(self)}>:\\nframe: \\n\" + str(self.frame) + \",\\nminimum distance:\" + str(self.min_dist)\n",
    "\n",
    "  @staticmethod\n",
    "  def cells_filter(t0:Cell,t1:Cell,disappeared_time:int,dist:float)->bool:\n",
    "    t0_f = None;\n",
    "    t1_f = None;\n",
    "    good_area = None;\n",
    "    try:\n",
    "      t0_f = t0.frame.iloc[0];\n",
    "      t1_f = t1.frame.iloc[0];\n",
    "      # raise Exception();  \n",
    "      good_area:bool = float(t1_f['area']) > (float(t0_f['area'])*(1-dfracsize)) and float(t1_f['area']) < (float(t0_f['area'])*(1+dfracsize));\n",
    "    except Exception as e:\n",
    "      print(t0,t1);\n",
    "      print(t0_f,t1_f);\n",
    "      print(t0_f['area'],t1_f['area']);\n",
    "      print(type(t0_f['area']),type(t1_f['area']));\n",
    "      print(good_area);\n",
    "      raise e;\n",
    "\n",
    "    if disappeared_time == 0:\n",
    "        result = dist < max_tracked_speed;\n",
    "    else:\n",
    "        if untracked_speed_type == SpeedType.linear:\n",
    "          result = dist < max_untracked_speed*disappeared_time;\n",
    "        elif untracked_speed_type == SpeedType.MSD:\n",
    "          result = dist**2 < max_untracked_diffusivity*disappeared_time;\n",
    "        else:\n",
    "          raise Exception();\n",
    "\n",
    "    if not(good_area) and result:\n",
    "        result = False;\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "frames = reading_cells_parameters['frames'];\n",
    "movies = reading_cells_parameters['movies'];\n",
    "m = master_bar(movies);\n",
    "for s in m:\n",
    "    print(f\"tracking movie #{s}\");\n",
    "    tracklists:Dict[int,list[pd.DataFrame]] = {} #objectid, frames\n",
    "    tracker = CentroidTracker[Cell](\n",
    "        Cell.cell_distance,\n",
    "        frame_filter=Cell.cells_filter,\n",
    "        maxDisappearedFrames=max_track_disappeared_time,\n",
    "        minPersistenceFrames=min_track_persistence_time);\n",
    "    for fr in progress_bar(frames[s],parent=m):\n",
    "\n",
    "        #get cells in current frame\n",
    "        maskcells:pd.DataFrame = sample_cells[(sample_cells['movie']==s) & (sample_cells['frame']==fr)]\n",
    "\n",
    "        cells = [Cell(maskcells.iloc[[i]]) for i in range(len(maskcells))];\n",
    "        objects = tracker.update(cells,allow_new = (fr < cell_collection_time or not(do_cell_collection)));        \n",
    "        pass;\n",
    "        for id,cell in objects.items():\n",
    "            frame = cell.frame.copy();\n",
    "            if id in tracklists:\n",
    "              if tracker.disappeared[id] > 0:\n",
    "                frame['label'] = -1; #disconnect repeated frame from mask label\n",
    "                frame['frame'] = fr;\n",
    "              tracklists[id].append(frame);\n",
    "            else:\n",
    "              tracklists[id] = [frame];\n",
    "    \n",
    "    #concat all tracks together\n",
    "    tracks:Dict[int,pd.DataFrame] = {tid:pd.concat(tracklists[tid],ignore_index=True) for tid in tracklists}\n",
    "    sample[s] = tracks\n",
    "\n",
    "with open(raw_tracks_path, 'wb') as handle:\n",
    "    pickle.dump(sample, handle)\n",
    "\n",
    "#transfer tracks and tracking parameters to folder_output in GCP bucket\n",
    "push_analysis();\n",
    "\n",
    "print(\"Tracking complete\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ivGlkgfsdc8H",
    "tags": []
   },
   "source": [
    "# Tracking Quality Control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KamiFxNiuzeM",
    "tags": []
   },
   "source": [
    "## Get analysis folder data from GCP bucket. Read tracks and cell shape features (centers will be used for display)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EluxmEzJZSVw",
    "tags": []
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "\n",
    "with open(cell_reading_params_path,'rb') as handle:\n",
    "    reading_cells_parameters = pickle.load(handle,encoding='latin1');\n",
    "\n",
    "with open(track_params_path,'rb') as handle:\n",
    "    track_params = pickle.load(handle,encoding='latin1');\n",
    "centertype = track_params['centroidtype'];\n",
    "    \n",
    "in_tracks = pd.read_pickle(raw_tracks_path);\n",
    "to_qc = in_tracks\n",
    "qc_params = {}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bVPzDTKrZRSg"
   },
   "source": [
    "## [Optional] Track Input Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Ensure tracking output is clean and orderly\n",
    "\n",
    "#@markdown basic validation: ensuring movies and frames match track values. Required for more complex validation. Highly recommended as a bare minimum\n",
    "do_validation = True #@param{type:\"boolean\"}\n",
    "\n",
    "#@markdown whether to enforce no gaps or repeat frames (necessary for track spacing and smoothing)\n",
    "validate_continuity = True #@param{type:\"boolean\"}\n",
    "\n",
    "\n",
    "qc_params.update({\"do_validation\":do_validation,\"valid_continuity\":validate_continuity})\n",
    "\n",
    "if do_validation:\n",
    "    valid_tracks:Dict[int,Dict[int,pd.DataFrame]] = {}\n",
    "    movies,frames = reading_cells_parameters['movies'],reading_cells_parameters['frames']\n",
    "    # print(movies,frames)\n",
    "    for movie,tracks in (m := master_bar(in_tracks.items())):\n",
    "        valid_tracks[movie] = {}\n",
    "        for tid,track in progress_bar(tracks.items(),parent=m):\n",
    "            \n",
    "            ##basic validation\n",
    "            track = track.sort_values('frame');\n",
    "            \n",
    "            vframes = track['frame'].isin(frames[movie])\n",
    "            if not vframes.all():\n",
    "                missing = list(track['frame'][vframes == False])\n",
    "                #seems unlikely to ever come up, but good to check regardless\n",
    "                raise Exception(f\"Track {tid} in movie {movie} has frames not in the official frame list at frames(s) {missing}\")\n",
    "            \n",
    "            vmovs = track['movie'].eq(movie)\n",
    "            if not vmovs.all():\n",
    "                missing = list(track['frame'][vmovs == False])\n",
    "                raise Exception(f\"Track {tid} in movie {movie} has incorrect movie labels at frame(s) {missing}\")\n",
    "                \n",
    "            ##validating continuity\n",
    "            if validate_continuity:\n",
    "                d = np.diff(track['frame'])\n",
    "                if (d != 1).any():\n",
    "                    # print(d)\n",
    "                    # print(np.argwhere(d != 1))\n",
    "                    # print(track['frame'])\n",
    "                    gaps = track['frame'].loc[np.squeeze(np.argwhere(d != 1))]\n",
    "                    raise Exception(f\"Track {tid} in movie {movie} contains nonconsecutive entries around frame(s) {gaps}\")\n",
    "\n",
    "            valid_tracks[movie][tid] = track;        \n",
    "    print(\"tracks validated successfully!\")\n",
    "    print(\"--basic validation checks passed\")\n",
    "    if validate_continuity:\n",
    "        print(\"--no discontinuous/repeated frames\")\n",
    "    to_qc = valid_tracks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## [Optional] Track Spacing and Smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Post Processing on Tracks [Requires Track Validation]\n",
    "\n",
    "#@markdown ## Spacing: Dropping track points to remove noise\n",
    "do_spacing = False #@param{type:\"boolean\"}\n",
    "\n",
    "#@markdown ### Spacing Gapsize: 1 is every other frame dropped, 2 is every three frames, 0 is no frames dropped\n",
    "spacing_gapsize = 1 #@param{type:\"integer\"}\n",
    "\n",
    "\n",
    "#@markdown ## Smoothing: Averaging between points to adjust for noise \n",
    "do_smoothing = True #@param{type:\"boolean\"}\n",
    "\n",
    "#@markdown ### Smoothing Width: how wide to do the rolling average. Width of 1 is no smoothing, with of two averages pairs of points, etc\n",
    "smoothing_width = 3 #@param{type:\"integer\"}\n",
    "\n",
    "#@markdown ### Smoothing Power: how many times to repeatedly perform the rolling average. Higher powers are generally less destructive to overall shape than wider averages\n",
    "smoothing_power = 3 #@param{type:\"integer\"}\n",
    "\n",
    "if do_spacing or do_smoothing:\n",
    "    \n",
    "    if not (qc_params['do_validation'] and qc_params['valid_continuity'] and \"valid_tracks\" in globals()):\n",
    "        raise Exception(\"Spacing and Smoothing requires tracks validated for continuity. Please enable continuity validation and run the above cell.\")\n",
    "\n",
    "    spsm_tracks:Dict[int,Dict[int,pd.DataFrame]] = {}\n",
    "\n",
    "    cx,cy = centertype + 'x', centertype + 'y'\n",
    "\n",
    "    ### guarantees single sets of consecutive, nonrepeating tracks\n",
    "    for mov,tracks in (m := master_bar(valid_tracks.items())):\n",
    "        spsm_tracks[mov] = {}\n",
    "        for tid,track in progress_bar(tracks.items(),parent=m):\n",
    "            track = track.sort_values('frame');\n",
    "            if do_spacing:\n",
    "                track = track[::spacing_gapsize+1]\n",
    "            if do_smoothing:\n",
    "                track[cx] = (moving_average**smoothing_power)(track[cx],smoothing_width,include_edges=True)\n",
    "                track[cy] = (moving_average**smoothing_power)(track[cy],smoothing_width,include_edges=True)\n",
    "            spsm_tracks[mov][tid] = track\n",
    "\n",
    "\n",
    "\n",
    "    to_qc = spsm_tracks\n",
    "\n",
    "tparams = {}\n",
    "    \n",
    "tparams['do_spacing'] = do_spacing;\n",
    "if do_spacing:\n",
    "    tparams['spacing_gapsize'] = spacing_gapsize;\n",
    "\n",
    "tparams['do_smoothing'] = do_smoothing\n",
    "if do_smoothing:\n",
    "    tparams['smoothing_width'] = smoothing_width\n",
    "    tparams['smoothing_power'] = smoothing_power\n",
    "\n",
    "tparams['smoothing_type'] = [['raw','smoothing'],\n",
    "                             ['spacing','spacing_smoothing']][do_spacing][do_smoothing]\n",
    "    \n",
    "qc_params.update(tparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bVPzDTKrZRSg",
    "tags": []
   },
   "source": [
    "## Perform QC operations, push parameters to GCP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yUHnB9p6dc8Q"
   },
   "source": [
    "keep: {movie:[track1,track2,etc]} - note that any unspecified movies will be left with all tracks\n",
    "\n",
    "example:\n",
    "keep={4:[1],6:[1],8:[7],11:[1,2],14:[9],17:[4],18:[1],19:[1],23:[3],24:[1,2,6],28:[1]}\n",
    "\n",
    "trim: {(movie,track):(firstframe_keep,lastframe_keep)}; if track is -1, will trim the whole movie to that trim\n",
    "\n",
    "example:\n",
    "trim={(4,1):(1,6),(17,4):(1,10),(25,1):(1,31),(26,1):(1,6)}\n",
    "\n",
    "removemov: [movie]\n",
    "\n",
    "example:\n",
    "removemov=[1,5,10,12,22,27,30]\n",
    "\n",
    "exclude: [movie,track]\n",
    "\n",
    "exclude=[(3,7),(13,2)]\n",
    "\n",
    "Here \"movie\" is the number assigned by the microscope, in filename [basename]_s[movie]_t[frame].tif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pQBFZKVpZSV3",
    "outputId": "8b23aee2-d828-4c07-8bd8-a57f8b0dc414",
    "tags": []
   },
   "outputs": [],
   "source": [
    "load_libraries()\n",
    "\n",
    "#INPUT QC OPERATIONS \n",
    "#@markdown the minimum track length to include (frames)\n",
    "minTrackLength=60 #@param {type:\"raw\"}\n",
    "\n",
    "#@markdown the minimum displacement over the length of the track (pixels)\n",
    "minTrackDisplacement=65 #@param {type:\"raw\"}\n",
    "\n",
    "#@markdown how long to wait after a track's appearance before including it (buffer period for cell division/collisions, segmentation issues, etc)\n",
    "initialTrackDelay = 15 #@param {type:\"raw\"}\n",
    "\n",
    "if qc_params.get(\"do_spacing\",False): #normalize to new track time\n",
    "    minTrackLength //= (qc_params['spacing_gapsize']+1)\n",
    "    initialTrackDelay //= (qc_params['spacing_gapsize']+1)\n",
    "\n",
    "\n",
    "#@markdown dict of {movie:[track1,track2,etc]} specific tracks to keep from particular movies; *only* the tracks specified in the movie will be kept\n",
    "keep={} #@param {type:\"raw\"}\n",
    "\n",
    "#@markdown dict of {(movie,track):(startframe,endframe)} sets the frame bounds of specific tracks in the sample\n",
    "trim={}  #@param {type:\"raw\"}\n",
    "\n",
    "#@markdown list of specific movies to exclude (overrides keep)\n",
    "removemov=[]  #@param {type:\"raw\"}\n",
    "\n",
    "#@markdown list of (movie,track): specific tracks to exclude\n",
    "exclude=[]  #@param {type:\"raw\"}\n",
    "\n",
    "print(\"starting qc...\")\n",
    "#apply QC operations\n",
    "sampTrStatus, sample = apply_qc(to_qc,qc_log_path,centertype,\n",
    "                                keep,\n",
    "                                trim,\n",
    "                                removemov,\n",
    "                                exclude,\n",
    "                                minTrackLength=minTrackLength,\n",
    "                                minTrackDisplacement=minTrackDisplacement,\n",
    "                                initialTrackDelay=initialTrackDelay);\n",
    "##SampTrStatus: dict of {movie#:statuses}, where statuses is a dict of {trackid:status}; in this case, status = 0 means bad track, status = 1 means good track\n",
    "\n",
    "trackingChanged = True;\n",
    "\n",
    "qc_params.update({'tracks_status':sampTrStatus, 'qc_tracks':sample, \n",
    "            'minTrackLength':minTrackLength,\n",
    "            'minTrackDisplacement':minTrackDisplacement,\n",
    "            'initialTrackDelay': initialTrackDelay,\n",
    "            'keep':keep, 'trim':trim, 'removemov':removemov, \n",
    "            'exclude':exclude})\n",
    "\n",
    "with open(qc_params_path, 'wb') as handle:\n",
    "    pickle.dump(qc_params, handle)\n",
    "\n",
    "tracksreg={}\n",
    "for imov in sample:\n",
    "  tracksreg[imov] = {}\n",
    "  for itr in sample[imov]:\n",
    "    #if the track satus == 1 (track passed QC)\n",
    "    if sampTrStatus[imov][itr]==1:\n",
    "      tracksreg[imov][itr] = (sample[imov][itr])\n",
    "\n",
    "with open(qc_tracks_path, 'wb') as f:\n",
    "  pickle.dump(tracksreg,f)\n",
    "\n",
    "\n",
    "\n",
    "push_analysis();\n",
    "print(\"track qc complete\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "klezaRyzdc8Q",
    "tags": []
   },
   "source": [
    "# [Optional] Visualize Tracks for track QC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_2n8DspAMSgk"
   },
   "source": [
    "**Tracking Quality Control Instructions:**\n",
    "\n",
    "Select which movie to view by changing the movie field in the cell below and re-running. If you see a track you want to remove or trim, you can edit the options in the cell above and re-run it to apply those edits. If you want to change the tracking parameters themselves, you can change the parameters of the final cell in the previous section (Get Cell Features (table), labeled masks, and track) and re-run it to re-do tracking, then run the previous and following cells.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 479
    },
    "id": "FxW2t7xoZSV8",
    "outputId": "bb1f5d2a-0b33-457f-9238-f40213561414",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown Copy images and masks from GCP to colab (may take a few minutes, zip folders in GCP to speed this up)\n",
    "with capture_output():\n",
    "    folderimages = fetch_images();\n",
    "    foldermasks = fetch_labeled_cellmasks();\n",
    "    foldernucmasks = fetch_labeled_nucmasks();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown Load qc tracks\n",
    "fetch_analysis()\n",
    "with open(cell_reading_params_path,'rb') as handle:\n",
    "    reading_cells_parameters = pickle.load(handle,encoding='latin1');\n",
    "with open(track_params_path,'rb') as handle:\n",
    "    track_params = pickle.load(handle,encoding='latin1');\n",
    "centertype = track_params['centroidtype'];\n",
    "with open(qc_params_path, 'rb') as handle:\n",
    "    qc_params = pickle.load(handle)\n",
    "sample = qc_params[\"qc_tracks\"] #get trimmed tracks with invalid movies included\n",
    "sampTrStatus = qc_params[\"tracks_status\"] #get status of qc'd tracks for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uWXMTMCJZSWA",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Input Tracking Display Parameters\n",
    "#@markdown Nucleus intensity for combined image (between 0 and 1)\n",
    "nucfrac=0.7 #@param {type:\"number\"}\n",
    "\n",
    "#@markdown Frame downsample rate (how often a frame is shown); increase this number for a faster processing time but less fine time-control\n",
    "frspace=1 #@param {type:\"integer\"}\n",
    "\n",
    "#@markdown Mask intensity of cells segmented but not tracked (between 0 and 1);\n",
    "untracked_intensity = 0.2 #@param {type:\"number\"}\n",
    "\n",
    "#@markdown Mask intensity of cells segmented and tracked but excluded from tracking by quality control (will be in red)\n",
    "excluded_intensity = 0.3 #@param {type:\"number\"}\n",
    "\n",
    "#@markdown how much to resize the image by (the more downscaled, the faster the playing but the less the resolution)\n",
    "downfrac = 0.5 #@param {type:\"number\"};\n",
    "\n",
    "#@markdown how large to plot the cell centers (pixels)\n",
    "centersize = 2 #@param {type:\"integer\"}\n",
    "\n",
    "#@markdown how large to draw the track id# (scaling from base size)\n",
    "textsize =  0.5#@param {type:\"number\"}\n",
    "\n",
    "#@markdown ### Track Display\n",
    "\n",
    "#@markdown how thick to draw the cell tracks\n",
    "trackwidth = 1#@param {type:\"integer\"}\n",
    "\n",
    "trackingChanged = True;\n",
    "\n",
    "\n",
    "images = {} if 'images' not in globals() or trackingChanged else images; #dict of movie: {True:[with_names],False:[without_names]}\n",
    "masks = {} if 'masks' not in globals() or trackingChanged else masks; #dict of movie {True:[with_centers],False:[without_centers]}\n",
    "tracks_images = {} if 'tracks' not in globals() or trackingChanged else tracks_images; #dict of movie:list[tracks]\n",
    "\n",
    "process_ready = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PCd5p6nkZSWC",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ##Prepare tracking display function (Required for movie display)\n",
    "\n",
    "try:\n",
    "    from google.colab.patches import cv2_imshow\n",
    "except:\n",
    "    cv2_imshow = cv2.imshow\n",
    "\n",
    "textfont = cv2.FONT_HERSHEY_SIMPLEX\n",
    "\n",
    "\n",
    "imagebasename = reading_cells_parameters['basename'];\n",
    "movies = reading_cells_parameters['movies'];\n",
    "\n",
    "def get_frame(images,masks,tracks_images,frame,showCenters=False,showNames=False,showTime=False,tracks_display=\"Neither\"):\n",
    "    image = images[showNames][frame].copy();\n",
    "    mask = masks[showCenters][frame].copy();\n",
    "    if tracks_display != \"Neither\":\n",
    "      tracks = tracks_images[frame];\n",
    "      alpha = np.where((tracks != 0).any(axis=2),1,0); #pixels set to 1 where there is at least one nonzero element of the rgb (not black)\n",
    "      if tracks_display in [\"Image side\",\"Both\"]:\n",
    "        overlay_image_alpha(image,tracks,0,0,alpha);\n",
    "\n",
    "      if tracks_display in [\"Mask side\",\"Both\"]:\n",
    "        overlay_image_alpha(mask,tracks,0,0,alpha);\n",
    "\n",
    "\n",
    "    combined = np.hstack((mask,image))\n",
    "    combined = rescale_intensity(combined,out_range=np.uint8).astype(np.uint8);\n",
    "    if showTime:\n",
    "      combined = cv2.putText(combined,str(frame),(0,combined.shape[0]),textfont,textsize,(255,255,255),1);\n",
    "    return combined\n",
    "\n",
    "def f(frame,showCenters=False,showNames=False,tracks_display=\"Neither\"):\n",
    "  c = get_frame(images,masks,tracks_images,frame,showCenters=showCenters,showNames=showNames);\n",
    "  l.set_data(c);\n",
    "  fig.canvas.draw() #use with %matplotlib notebook\n",
    "  display(fig) #use with %matplotlib inline\n",
    "\n",
    "\n",
    "def overlay_image_alpha(img, img_overlay, x, y, alpha_mask):\n",
    "    \"\"\"Overlay `img_overlay` onto `img` at (x, y) and blend using `alpha_mask`.\n",
    "\n",
    "    `alpha_mask` must have same HxW as `img_overlay` and values in range [0, 1].\n",
    "    \"\"\"\n",
    "    # Image ranges\n",
    "    y1, y2 = max(0, y), min(img.shape[0], y + img_overlay.shape[0])\n",
    "    x1, x2 = max(0, x), min(img.shape[1], x + img_overlay.shape[1])\n",
    "\n",
    "    # Overlay ranges\n",
    "    y1o, y2o = max(0, -y), min(img_overlay.shape[0], img.shape[0] - y)\n",
    "    x1o, x2o = max(0, -x), min(img_overlay.shape[1], img.shape[1] - x)\n",
    "\n",
    "    # Exit if nothing to do\n",
    "    if y1 >= y2 or x1 >= x2 or y1o >= y2o or x1o >= x2o:\n",
    "        return\n",
    "\n",
    "    # Blend overlay within the determined ranges\n",
    "    img_crop = img[y1:y2, x1:x2,:]\n",
    "    img_overlay_crop = img_overlay[y1o:y2o, x1o:x2o,:]\n",
    "    alpha = alpha_mask[y1o:y2o, x1o:x2o, np.newaxis]\n",
    "    alpha_inv = 1.0 - alpha\n",
    "\n",
    "    img_crop[:] = alpha * img_overlay_crop + alpha_inv * img_crop\n",
    "\n",
    "def process_movie(movie:int,frames:list,parentbar=None):\n",
    "      if movie not in movies:\n",
    "        raise Exception(f\"movie selection {movie} not in experimental list of movies {movies}\");\n",
    "      if not process_ready:\n",
    "        raise Exception(\"Attempted to process movies with out of date parameters; run the \\\"Prepare Tracking Function\\\" Cell\")\n",
    "      movie_tracks = sample[movie];\n",
    "      print(f\"processing movie #{movie}\");\n",
    "    \n",
    "      \n",
    "      #random color per id, will be consistent for that id within the movie\n",
    "      centerColors = DefaultDict(lambda: (random.randrange(0,256),random.randrange(0,256),random.randrange(0,256)));\n",
    "      grey = (50,50,50);\n",
    "      [centerColors.update([(id,grey)]) for id,status in sampTrStatus[movie].items() if not status];\n",
    "      \n",
    "      acc_tracks_image = None;\n",
    "      prev_track_pos = None;\n",
    "\n",
    "      timage = {True:{},False:{}}\n",
    "      tmask = {True:{},False:{}}\n",
    "      ttrack_images = {}\n",
    "\n",
    "      for i in progress_bar(frames,parent=(parentbar or None)):\n",
    "          filename_mask = imagebasename + \"_s\" + str(movie)+\"_t\"+str(i)+\".TIF\";\n",
    "\n",
    "          maskmem = imread(foldermasks/filename_mask);\n",
    "          masknuc = imread(foldernucmasks/filename_mask);\n",
    "          \n",
    "          #read image\n",
    "          filename_img = imagebasename + \"_s\" + str(movie)+\"_t\"+str(i)+\".tif\";\n",
    "          imageorig = imread(folderimages/filename_img)\n",
    "\n",
    "          #rescale image intensity\n",
    "          imageorig = rescale_intensity(imageorig);\n",
    "          image=resize(imageorig, (maskmem.shape[0] * downfrac , maskmem.shape[1] * downfrac),  anti_aliasing=True); #this order matters so that the data types work out\n",
    "          image = rescale_intensity(image,out_range=np.uint8).astype(np.uint8);\n",
    "          image = np.stack((image,image,image),axis=2); #make color image\n",
    "          \n",
    "          \n",
    "          \n",
    "          trackedLabels = [];\n",
    "          rejectedLabels = [];\n",
    "          centers = {};\n",
    "          for tid,data in movie_tracks.items():\n",
    "            fDat = data[data['frame']==i];\n",
    "            if not fDat.empty:\n",
    "              trackedLabels.append(fDat)\n",
    "              if not sampTrStatus[movie][tid]:\n",
    "                rejectedLabels.append(fDat);\n",
    "              centers[tid] = fDat[[centertype+'x',centertype+'y']].reset_index();\n",
    "          \n",
    "          \n",
    "          #get the label of every tracked mask in this frame\n",
    "          trackedLabels = pd.concat(trackedLabels)['label'] if len(trackedLabels) > 0 else [];\n",
    "          rejectedLabels = pd.concat(rejectedLabels)['label'] if len(rejectedLabels) > 0 else []\n",
    "\n",
    "\n",
    "          #create bitmasks of untracked and rejected cells\n",
    "          cellmask = (maskmem != 0);\n",
    "          untracked = np.isin(maskmem,trackedLabels,invert=True) & cellmask\n",
    "          rejected = np.isin(maskmem,rejectedLabels);\n",
    "          \n",
    "          #draw tracks\n",
    "          if acc_tracks_image is None:\n",
    "            acc_tracks_image = np.zeros(np.array(image.shape),dtype=np.uint8);\n",
    "            prev_track_pos = centers;\n",
    "          else:\n",
    "            for id in centers:\n",
    "              if id in prev_track_pos:\n",
    "                prev = prev_track_pos[id];\n",
    "                prev = (int(prev[centertype+'x'][0]*downfrac),int(prev[centertype+'y'][0]*downfrac));\n",
    "\n",
    "                pos = centers[id];\n",
    "                pos = (int(pos[centertype+'x'][0]*downfrac),int(pos[centertype+'y'][0]*downfrac));\n",
    "\n",
    "                acc_tracks_image = cv2.line(acc_tracks_image,prev,pos,centerColors[id],trackwidth);\n",
    "            prev_track_pos = centers;\n",
    "\n",
    "          ttrack_images[i] = acc_tracks_image.copy();        \n",
    "\n",
    "          #unlabel for visualization\n",
    "          maskmem[maskmem>=1] = 1\n",
    "\n",
    "          #combine membrane and nucleus masks\n",
    "          maskcomb = (maskmem - nucfrac*masknuc)\n",
    "          maskcomb[maskcomb<0]=0; #floating point stuff, this is so stupid\n",
    "\n",
    "          #apply untracked and excluded intensities\n",
    "          maskcomb[untracked] *= untracked_intensity;\n",
    "          maskcomb[rejected] *= excluded_intensity\n",
    "\n",
    "          #rescale to full intensity, make int8 again\n",
    "          maskcomb = rescale_intensity(maskcomb,out_range=np.uint8).astype(np.uint8);\n",
    "\n",
    "          #color rejected cells red\n",
    "          qcomb = (maskcomb//4).astype(maskcomb.dtype);\n",
    "          halfmask = np.where(rejected,qcomb,maskcomb);\n",
    "          \n",
    "          #turn into rgb image\n",
    "          maskcomb = np.stack((maskcomb,halfmask,halfmask),axis=2);\n",
    "\n",
    "          #downscale\n",
    "          maskcomb=resize(maskcomb,(int(maskcomb.shape[0]*downfrac), int(maskcomb.shape[1]*downfrac)),preserve_range=True).astype(np.uint8);\n",
    "\n",
    "          #save unannotated frames\n",
    "          tmask[False][i] = maskcomb.copy();\n",
    "          timage[False][i] = image.copy();\n",
    "\n",
    "          #annotate with centers and names\n",
    "          for id,pos in centers.items():\n",
    "            pos = (int(pos[centertype+'x'][0]*downfrac),int(pos[centertype+'y'][0]*downfrac));\n",
    "            maskcomb = cv2.circle(maskcomb,pos,2,centerColors[id],-1);\n",
    "            image = cv2.putText(image,str(id),pos,textfont,textsize,centerColors[id],1);\n",
    "\n",
    "          #save annotated frames\n",
    "          tmask[True][i] = maskcomb;\n",
    "          timage[True][i] = image;\n",
    "      return timage,tmask,ttrack_images\n",
    "\n",
    "process_ready = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dzGrXoJRZSWN",
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "#@markdown ### Compute Tracking Display Frames\n",
    "#@markdown Get combined images & masks for visualization, draw mask centers and labeled names\n",
    "\n",
    "#@markdown Which movie to load\n",
    "movie_selection=6 #@param {type:\"integer\"}\n",
    "\n",
    "#@markdown if you want to process the other movies while loading the selection\n",
    "process_all = False #@param {type:\"boolean\"}\n",
    "\n",
    "#@markdown Force reprocessing of selected movie (select if running this cell with new data/parameters says \"previously prepared)\n",
    "force_reprocess = False #@param {type:\"boolean\"}\n",
    "\n",
    "to_process = master_bar(movies) if process_all else [movie_selection];\n",
    "\n",
    "\n",
    "for movie in to_process:\n",
    "  frames = reading_cells_parameters['frames'][movie][::frspace];\n",
    "\n",
    "  movie_tracked = True;\n",
    "  if movie not in images or len(images[movie][True]) != len(frames):\n",
    "    movie_tracked = False;\n",
    "  if movie not in masks or len(masks[movie][True]) != len(frames):\n",
    "    movie_tracked = False;\n",
    "  if movie not in tracks_images or len(tracks_images) != len(frames):\n",
    "    movie_tracked = False;\n",
    "\n",
    "  if not movie_tracked or (movie == movie_selection and force_reprocess):\n",
    "    images[movie],masks[movie],tracks_images[movie] = process_movie(movie,frames,to_process if process_all else None);\n",
    "trackingChanged = False;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lp3v8pIxZSWP",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## View Tracking Output - Mediapy\n",
    "#@markdown ### (Requires computing display frames for selected movie)\n",
    "#@markdown #-\n",
    "#@markdown ### Mediapy results in a much cleaner video experience than matlab, but can't be changed dynamically\n",
    "#@markdown Framerate: speed of movie playing in frames per second (not limited by processing time, can go as fast as you want);\n",
    "framerate =  1#@param {type:\"number\"}\n",
    "#@markdown ### Annotation Parameters:\n",
    "centers = True #@param {type:\"boolean\"}\n",
    "names = True #@param {type:\"boolean\"}\n",
    "frameNumber = True #@param {type:\"boolean\"}\n",
    "#@markdown whether and where to draw the tracks of each cell on the image frames\n",
    "display_tracks_side = \"Image side\" #@param [\"Neither\", \"Mask side\", \"Image side\", \"Both\"]\n",
    "\n",
    "\n",
    "\n",
    "frlist = [[f] for f in reading_cells_parameters['frames'][movie][::frspace]];\n",
    "print(frlist)\n",
    "try:\n",
    "  media.show_video(starmap(partial(get_frame,images[movie_selection],masks[movie_selection],tracks_images[movie_selection],showCenters=centers,showNames=names,showTime=frameNumber,tracks_display=display_tracks_side),progress_bar(frlist)),fps=framerate);\n",
    "except NameError as e:\n",
    "  raise Exception(\"Unable to load movie - did you set up the frame function?\") from e;\n",
    "except KeyError as k:\n",
    "  raise Exception(\"Unable to load entire movie - run the parameters and computation cells to ensure full movie complete\") from k;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nmz7QriiZSWT",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Save Video Output (Mediapy)\n",
    "#@markdown ### (Requires computing display frames for selected movie)\n",
    "#@markdown #-\n",
    "#@markdown Framerate: speed of movie playing in frames per second (not limited by processing time, can go as fast as you want);\n",
    "framerate = 3#@param {type:\"number\"}\n",
    "#@markdown ### Annotation Parameters:\n",
    "centers = True #@param {type:\"boolean\"}\n",
    "names = True #@param {type:\"boolean\"}\n",
    "frameNumber = True #@param {type:\"boolean\"}\n",
    "#@markdown whether and where to draw the tracks of each cell on the image frames\n",
    "display_tracks_side = \"Mask side\" #@param [\"Neither\", \"Mask side\", \"Image side\", \"Both\"]\n",
    "\n",
    "#@markdown GCP output filename (to automatically name the path based on movie#, put \"{movie}\" in the name)\n",
    "gcp_out_path = \"trackqc_using_cell_collection_ex41_mov{movie}.mp4\" #@param {type:\"string\"}\n",
    "gcp_out_path:Path = Path(format_path(gcp_out_path,movie=movie_selection,collection=track_params['do_cell_collection']));\n",
    "\n",
    "#@markdown ##### Treat the above path as relative to the experiment's analysis output folder?\n",
    "in_analysis = True #@param {type:\"boolean\"}\n",
    "\n",
    "local_path = local_analysis_output_folder / gcp_out_path if in_analysis else temp_folder / Path(gcp_out_path).name\n",
    "gcp_out_path = gcp_analysis_output_folder / gcp_out_path if in_analysis else Path(gcp_out_path);\n",
    "frlist = [[f] for f in reading_cells_parameters['frames'][movie][::frspace]];\n",
    "try:\n",
    "  media.write_video(local_path,starmap(partial(get_frame,images[movie_selection],masks[movie_selection],tracks_images[movie_selection],showCenters=centers,showNames=names,showTime=frameNumber,tracks_display=display_tracks_side),progress_bar(frlist)),fps=framerate);\n",
    "except NameError as e:\n",
    "  raise Exception(\"Unable to load movie - did you set up the frame function?\") from e;\n",
    "except KeyError as k:\n",
    "  raise Exception(\"Unable to load entire movie - run the parameters and computation cells to ensure full movie complete\") from k;\n",
    "\n",
    "!gsutil -m cp \"{gs_str(local_path)}\" \"{gs_str(gcp_out_path)}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3dVhi9ZTZSWX",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Process and save range of movies\n",
    "#@markdown ### Requires process function initialization, but not computing display frames for the selected movies\n",
    "#@markdown #-\n",
    "#@markdown ## Doing it in batches saves interaction time and doesn't remember results, freeing up fram usage\n",
    "#@markdown If movie is -1, will include all movies\n",
    "movie = -1#@param{type:\"integer\"}\n",
    "\n",
    "#@markdown If you want to include a range of movies, setting this to a positive number will process all movies from movie to movie_end\n",
    "movie_end =  12#@param{type:\"integer\"}\n",
    "\n",
    "selected_movies = [];\n",
    "if movie == -1:\n",
    "    selected_movies = sample.keys()\n",
    "elif movie_end > 0:\n",
    "    selected_movies = [m for m in range(movie,movie_end+1) if m in sample]\n",
    "elif (movie not in sample):\n",
    "   raise Exception(f\"Error: movie #{movie} not in experimental movies {list(sample.keys())}\");\n",
    "else:\n",
    "    selected_movies = [movie];\n",
    "\n",
    "print(\"processing movies\",list(selected_movies))\n",
    "\n",
    "#@markdown Framespace: how often a frame is included (1 is every frame, 2 is every other frame, etc)\n",
    "framespace = 1#@param {type:\"integer\"}\n",
    "#@markdown Framerate: speed of movie playing in frames per second (not limited by processing time, can go as fast as you want);\n",
    "framerate = 3#@param {type:\"number\"}\n",
    "#@markdown ### Annotation Parameters:\n",
    "centers = True #@param {type:\"boolean\"}\n",
    "names = True #@param {type:\"boolean\"}\n",
    "frameNumber = True #@param {type:\"boolean\"}\n",
    "#@markdown whether and where to draw the tracks of each cell on the image frames\n",
    "display_tracks_side = \"Mask side\" #@param [\"Neither\", \"Mask side\", \"Image side\", \"Both\"]\n",
    "\n",
    "#@markdown GCP output filename (Supported formatting: {collection}, {movie})\n",
    "gcp_out_str = \"trackqc_mov{movie}.mp4\" #@param {type:\"string\"}\n",
    "#@markdown ##### Treat the above path as relative to the experiment's analysis output folder?\n",
    "in_analysis = True #@param {type:\"boolean\"}\n",
    "bar = master_bar(selected_movies)\n",
    "for movie in bar:\n",
    "  \n",
    "  gcp_out_path = Path(format_path(gcp_out_str,movie=movie,collection=track_params['do_cell_collection'],qc_params=qc_params));\n",
    "\n",
    "  local_path = local_analysis_output_folder / gcp_out_path if in_analysis else temp_folder / Path(gcp_out_path).name\n",
    "  gcp_out_path = gcp_analysis_output_folder / gcp_out_path if in_analysis else Path(gcp_out_path);\n",
    "\n",
    "  frames = reading_cells_parameters['frames'][movie][::framespace];\n",
    "\n",
    "  # print(frames);\n",
    "  args = process_movie(movie,frames,parentbar=bar);\n",
    "\n",
    "  try:\n",
    "    media.write_video(local_path,map(partial(get_frame,*args,showCenters=centers,showTime=frameNumber,showNames=names,tracks_display=display_tracks_side),progress_bar(frames)),fps=framerate,crf=25); #lower crf is higher quality\n",
    "  except NameError as e:\n",
    "    raise Exception(\"Unable to load movie - did you set up the frame function?\") from e;\n",
    "  except KeyError as k:\n",
    "    raise Exception(\"Unable to load entire movie - run the parameters and computation cells to ensure full movie complete\") from k;\n",
    "\n",
    "  del args\n",
    "\n",
    "  !gsutil -m cp \"{gs_str(local_path)}\" \"{gs_str(gcp_out_path)}\"\n",
    "    \n",
    "push_analysis();\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xwR0SQhEZSWa"
   },
   "outputs": [],
   "source": [
    "#@markdown ## View Tracking Output - Matplotlib\n",
    "#@markdown ### Slower and rougher playing experience, but dynamically changeable\n",
    "import time\n",
    "import ipywidgets\n",
    "#@markdown Framerate: how quickly movie will be shown in frames per second (will be limited by image display speed)\n",
    "framerate =  2#@param {\"type\":\"number\"}\n",
    "delay = 1/framerate;\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(14, 8))\n",
    "ax_mask = fig.add_subplot(111) \n",
    "l = ax_mask.imshow(get_frame(frames[0]));\n",
    "\n",
    "\n",
    "play = widgets.Play(\n",
    "    min=frames[0],\n",
    "    max=frames[-1],\n",
    "    step=frspace,\n",
    "    value=frames[0],\n",
    "    interval=delay*1000,\n",
    "    description=\"Press play\",\n",
    "    disabled=False\n",
    ")\n",
    "slider = widgets.IntSlider(msg_throttle=1)\n",
    "widgets.jslink((play, 'value'), (slider, 'value'))\n",
    "# display(widgets.HBox([play, slider]))\n",
    "display(play);\n",
    "interact(f,frame=slider,showCenters=False,showNames=False);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BO5N4SrLHfSx",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Save masks for tracks in tiff stacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JEP9nO4QZSWf"
   },
   "outputs": [],
   "source": [
    "#function to read a mask corresponding to a given movie (implicit parameter), and frame (explicit parameter iframe)\n",
    "#get_mask = lambda iframe: imread( masks_folder/basename+'_s'+str(int(movie))+'_t'+str(int(frame[iframe]))+'.TIF')\n",
    "def get_mask(iframe):\n",
    "  filename = basename+'_s'+str(int(movie))+'_t'+str(int(frame[iframe]))+'.TIF'\n",
    "  return imread(masks_folder/filename)\n",
    "\n",
    "#function that returns the mask of a cell centered. It takes the cell label in the mask, and the movie\n",
    "#as implicit parameters determined by the track information (as used later). \n",
    "def get_centered_cell(iframe):\n",
    "  mask=get_mask(iframe)\n",
    "  #erase all objects with different label as cell\n",
    "  mask[mask != label[iframe] ] =0\n",
    "  #set cell positions as 1\n",
    "  mask[mask > 0 ] = 1\n",
    "  props = measure.regionprops(mask)\n",
    "  #center image\n",
    "  centroid=props[0].centroid\n",
    "  N=mask.shape\n",
    "  centered=mask[np.ix_((np.arange(N[0]) + int(centroid[0]) - int(N[0]/2)) % N[0] , (np.arange(N[1]) + int(centroid[1]) - int(N[1]/2)) % N[1])]\n",
    "  return centered\n",
    "\n",
    "#function that returns the mask of a cell. It takes the cell label in the mask, and the movie\n",
    "#as implicit parameters determined by the track information (as used later).\n",
    "def get_cell(iframe):\n",
    "  mask=get_mask(iframe)  \n",
    "  #erase all objects with different label as cell\n",
    "  mask[mask != label[iframe] ] =0\n",
    "  #set cell positions as 1\n",
    "  mask[mask > 0 ] = 1\n",
    "\n",
    "  return mask\n",
    "\n",
    "rotation_matrix = lambda angle: np.asarray([[np.cos(angle) , -np.sin(angle)],[np.sin(angle), np.cos(angle)]] )\n",
    "\n",
    "\n",
    "def skew_from_hist(hist):\n",
    "  value = np.asarray(range(len(hist)))\n",
    "  mean = np.sum(hist*value)/np.sum(hist)\n",
    "  m2 = np.sum(hist*(value-mean)**2)/np.sum(hist)\n",
    "  m3 = np.sum(hist*(value-mean)**3)/np.sum(hist)\n",
    "  return  m3/m2**(3/2) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xB6XU-PqZSWj"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "\n",
    "\n",
    "#unzip labeld masks\n",
    "#!unzip {local_analysis_output_folder}/labeledmasks.zip\n",
    "\n",
    "masks_folder = fetch_labeled_cellmasks()\n",
    "\n",
    "#read pickled tracks\n",
    "#unpickle tracks  \n",
    "with open(qc_tracks_path, 'rb') as handle:\n",
    "  tracks_noshape = pickle.load(handle, encoding='latin1') \n",
    "\n",
    "#create folder to store masks from tracks\n",
    "if not os.path.exists(temp_folder/'tracks_masks'):\n",
    "  os.mkdir(temp_folder/'tracks_masks');\n",
    "#!mkdir temp_folder/tracks_masks\n",
    "\n",
    "\n",
    "###Loop to add features to each timepoint of all the tracks\n",
    "#Add morphological and morphodynamical features to each timepoint (frame) in the tracks\n",
    "\n",
    "tracks = copy.copy(tracks_noshape)\n",
    "\n",
    "for movie in tracks:\n",
    "  for itrack in tracks[movie]:\n",
    "    tracklength = len(tracks[movie][itrack])\n",
    "    frame = list(tracks[movie][itrack]['frame'])\n",
    "    label = list(tracks[movie][itrack]['label'])\n",
    "    masks=os.listdir(masks_folder)\n",
    "    #implicit argument to read cells\n",
    "    basename=re.findall(r\"(.+)_s\",masks[0])[0] \n",
    "\n",
    "    track_masks=[]\n",
    "\n",
    "\n",
    "    #GET SCIKIT-IMAGE CELL METRICS AND SAVE CELL MASK\n",
    "    for iframe in range(tracklength):\n",
    "      cell = get_cell(iframe) \n",
    "\n",
    "      #store cell mask\n",
    "      track_masks.append(cell)\n",
    "\n",
    "    #SAVE TRACK MASKS AS TIFF STACKS\n",
    "    track_masks = np.asarray(track_masks)\n",
    "    savepath = \"tracks_masks/\" + experiment + \"_movie\" + str(movie) + \"_track\"+ str(itrack) + '.TIF'\n",
    "    imsave(temp_folder/savepath , track_masks, check_contrast=False)\n",
    "\n",
    "#transfer data (tracks_masks and tracks_shape) to GCP bucket\n",
    "!zip -r {local_analysis_output_folder}/tracks_masks.zip {temp_folder}/tracks_masks\n",
    "\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YThS6UbsGy4u",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Shape2tracks - tracks with shape info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iiukRaOvZSWm"
   },
   "outputs": [],
   "source": [
    "#@markdown label of substrate or region on gel where experiment was taken\n",
    "\n",
    "region = 'gel' #@param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IXbxaSm5ZSWo"
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "\n",
    "\n",
    "#unzip labeld masks\n",
    "#!unzip {local_analysis_output_folder}/labeledmasks.zip\n",
    "\n",
    "masks_folder = fetch_labeled_cellmasks()\n",
    "\n",
    "#read pickled tracks\n",
    "#unpickle tracks  \n",
    "with open(qc_tracks_path, 'rb') as handle:\n",
    "  tracks_noshape = pickle.load(handle, encoding='latin1') \n",
    "\n",
    "\n",
    "###Loop to add features to each timepoint of all the tracks\n",
    "#Add morphological and morphodynamical features to each timepoint (frame) in the tracks\n",
    "\n",
    "tracks = copy.copy(tracks_noshape)\n",
    "\n",
    "for movie in tracks:\n",
    "  for itrack in tracks[movie]:\n",
    "    tracklength = len(tracks[movie][itrack])\n",
    "    #store experiment and id_track, to retrieve track_masks from track\n",
    "    tracks[movie][itrack][\"experiment\"]=experiment\n",
    "    tracks[movie][itrack][\"region\"]=region\n",
    "    tracks[movie][itrack][\"track_id\"]=itrack\n",
    "    #read data to retrieve corresponding labeled mask cell\n",
    "    #movie = tracks[itrack]['movie'].iloc[0]\n",
    "    frame = list(tracks[movie][itrack]['frame'])\n",
    "    label = list(tracks[movie][itrack]['label'])\n",
    "    masks=os.listdir(masks_folder)\n",
    "    #implicit argument to read cells\n",
    "    basename=re.findall(r\"(.+)_s\",masks[0])[0] \n",
    "\n",
    "    #get polarization angle, skew, protrusion and retraction angle,norm area, radii \n",
    "\n",
    "    median_centroidy = []\n",
    "    median_centroidx = []\n",
    "    protrusion_angles=[]\n",
    "    mean_protrusion_angles=[]\n",
    "    protrusion_norm_radii =[]\n",
    "    protrusion_norm_areas=[]\n",
    "    retraction_angles=[]\n",
    "    mean_retraction_angles=[]\n",
    "    retraction_norm_areas=[]\n",
    "    retraction_norm_radii=[]\n",
    "    mean_retraction_norm_radii=[]\n",
    "    mean_protrusion_norm_radii =[]\n",
    "\n",
    "    cell_angles = []\n",
    "    cell_skews = []\n",
    "\n",
    "    #select scikit-image shape metrics\n",
    "    shape_metrics = ['area','convex_area','eccentricity','orientation','perimeter','equivalent_diameter','solidity','extent','major_axis_length','minor_axis_length','centroid']\n",
    "\n",
    "    track_shape_metrics = pd.DataFrame()\n",
    "\n",
    "    #GET SCIKIT-IMAGE CELL METRICS AND SAVE CELL MASK\n",
    "    for iframe in range(tracklength):\n",
    "      cell = get_cell(iframe) \n",
    "      if np.any(cell) == False:\n",
    "        cell_shape_metrics = measure.regionprops_table(cell, properties = shape_metrics)        \n",
    "        cell_shape_metrics=pd.DataFrame(cell_shape_metrics)\n",
    "        row = np.empty(len(shape_metrics)+1)\n",
    "        row[:] = np.nan\n",
    "        cell_shape_metrics = cell_shape_metrics.append(pd.Series(row, index=cell_shape_metrics.columns),ignore_index=True)\n",
    "        #append cell shape metrics to track shape metrics\n",
    "        track_shape_metrics = track_shape_metrics.append(cell_shape_metrics, ignore_index=True)\n",
    "        #get median calculated centroid\n",
    "        celly = np.array(np.nan)\n",
    "        cellx = np.array(np.nan)\n",
    "        median_centroidy.append( np.median(celly) ) \n",
    "        median_centroidx.append( np.median(cellx) )\n",
    "      else:\n",
    "        cell_shape_metrics = measure.regionprops_table(cell, properties = shape_metrics)        \n",
    "        cell_shape_metrics=pd.DataFrame(cell_shape_metrics)\n",
    "        #append cell shape metrics to track shape metrics\n",
    "        track_shape_metrics = track_shape_metrics.append(cell_shape_metrics, ignore_index=True)\n",
    "        #get median calculated centroid\n",
    "        celly,cellx = np.where(cell)\n",
    "        median_centroidy.append( np.median(celly) ) \n",
    "        median_centroidx.append( np.median(cellx) )\n",
    "\n",
    "\n",
    "\n",
    "    for iframe in range(tracklength-1):\n",
    "\n",
    "      #GET CELL POLARIZATION ANGLE \n",
    "      cell = get_cell(iframe)\n",
    "      if np.any(cell) == False:\n",
    "        cell_angles.append( np.nan )\n",
    "        cell_skews.append(np.nan)\n",
    "      else:\n",
    "        cell_centered = get_centered_cell(iframe)\n",
    "        props = regionprops(cell_centered)\n",
    "        #major axis angle with respect to y (0 axis : rows) counter-clockwise -pi/2 , pi/2  \n",
    "        angle_y = props[0].orientation\n",
    "        #rotate cell so that major axis is aligned with the y axis\n",
    "        cell_y = rotate(cell_centered,-angle_y*180/np.pi,order=0, preserve_range=True)\n",
    "        #project cell on the x (1) axis\n",
    "        proj_x = np.sum(cell_y,1)      \n",
    "        #get cell polarization vector in a regular coordinate system\n",
    "        #skew sign corresponds to the tail of the distribution, cell polarization is  \n",
    "        #defined here as -skew. Because image y-axis is inverted, in a regular coordinate\n",
    "        #system, cell polarization vector is defined as -(-skew) = skew:\n",
    "        skew_cell = skew_from_hist(proj_x)\n",
    "        celly_polarization = [0 , skew_cell ]      \n",
    "        #rotate back\n",
    "        polarization = np.dot( rotation_matrix(angle_y) , celly_polarization )\n",
    "        cell_angles.append( np.arctan2(polarization[1],polarization[0]) )\n",
    "        cell_skews.append(abs(skew_cell)) \n",
    "\n",
    "\n",
    "      #PROTRUSION AND RETRACTION VECTORS : angle, norm_areas, radii (calculated with medians)\n",
    "      if np.any(cell) == False or np.any(get_cell(iframe+1)) == False:\n",
    "        protrusion_angles.append( np.nan )\n",
    "        protrusion_norm_areas.append(np.nan)\n",
    "        #effective radius (area/pi)^0.5\n",
    "        protrusion_norm_radii.append( np.nan )\n",
    "        retraction_angles.append(np.nan)\n",
    "        retraction_norm_areas.append( np.nan )\n",
    "        retraction_norm_radii.append( np.nan )\n",
    "        mean_protrusion_angles.append(np.nan)\n",
    "        mean_retraction_angles.append(np.nan )\n",
    "        mean_protrusion_norm_radii.append( np.nan )\n",
    "        mean_retraction_norm_radii.append( np.nan )\n",
    "\n",
    "        med_centroids = pd.DataFrame({'median_centroidx':[np.nan], 'median_centroidy':[np.nan]})\n",
    "\n",
    "      else:\n",
    "        difference  = get_centered_cell(iframe+1) - get_centered_cell(iframe)\n",
    "        #get centroids of cell(iframe), protrusion and retraction\n",
    "        y,x =np.where(get_centered_cell(iframe))\n",
    "        centroidy, centroidx = np.median(y), np.median(x)\n",
    "        yp,xp = np.where(difference==1)\n",
    "        protrusion_y, protrusion_x  = np.median(yp) , np.median(xp)\n",
    "        yr,xr = np.where(difference==-1)\n",
    "        retraction_y, retraction_x  = np.median(yr) , np.median(xr) \n",
    "        #get protr and retr angle in a regular coordinate system: [ximage , - yimage]\n",
    "        #and norm_areas\n",
    "        protrusion_angles.append( np.arctan2( -(protrusion_y - centroidy) , protrusion_x - centroidx) )\n",
    "        protrusion_norm_areas.append(len(yp)/ len(y))\n",
    "        #effective radius (area/pi)^0.5\n",
    "        protrusion_norm_radii.append( ((protrusion_x-centroidx)**2+(protrusion_y-centroidy)**2)**0.5/(len(y)/np.pi)**0.5 )\n",
    "        retraction_angles.append(np.arctan2( -(retraction_y - centroidy), retraction_x - centroidx)  )\n",
    "        retraction_norm_areas.append( len(yr)/len(y) )\n",
    "        retraction_norm_radii.append( ((retraction_x-centroidx)**2+(retraction_y-centroidy)**2)**0.5/(len(y)/np.pi)**0.5 )\n",
    "\n",
    "\n",
    "        #PROTRUSION AND RETRACTION angles, radii, calculated with means\n",
    "        #get centroids of cell(iframe), protrusion and retraction\n",
    "        mean_centroidy, mean_centroidx = np.mean(y), np.mean(x)\n",
    "        mean_protrusion_y, mean_protrusion_x  = np.mean(yp) , np.mean(xp)\n",
    "        mean_retraction_y, mean_retraction_x  = np.mean(yr) , np.mean(xr) \n",
    "        #get protr and retr angle in a regular coordinate system: [ximage , - yimage]\n",
    "        #and norm_areas\n",
    "        mean_protrusion_angles.append(np.arctan2( -(mean_protrusion_y - mean_centroidy), mean_protrusion_x - mean_centroidx) )\n",
    "        mean_retraction_angles.append(np.arctan2( -(mean_retraction_y - mean_centroidy), mean_retraction_x - mean_centroidx) )\n",
    "        mean_protrusion_norm_radii.append( ((mean_protrusion_x-mean_centroidx)**2+(mean_protrusion_y-mean_centroidy)**2)**0.5/(len(y)/np.pi)**0.5 )\n",
    "        mean_retraction_norm_radii.append( ((mean_retraction_x-mean_centroidx)**2+(mean_retraction_y-mean_centroidy)**2)**0.5/(len(y)/np.pi)**0.5 )\n",
    "\n",
    "\n",
    "        med_centroids = pd.DataFrame({'median_centroidx':median_centroidx, 'median_centroidy':median_centroidy})\n",
    "\n",
    "    shape_features= pd.DataFrame({'polarity_angle':cell_angles, 'abs-skew':cell_skews, \n",
    "                              'protr_angle':protrusion_angles, 'mean_protr_angle':mean_protrusion_angles, 'protr_norm_area':protrusion_norm_areas, \n",
    "                          'retr_angle': retraction_angles, 'mean_retr_angle': mean_retraction_angles, 'retr_norm_area': retraction_norm_areas,\n",
    "                          'protr_norm_radii':protrusion_norm_radii, 'mean_protr_norm_radii':mean_protrusion_norm_radii,\n",
    "                          'retr_norm_radii':retraction_norm_radii, 'mean_retr_norm_radii':mean_retraction_norm_radii, })\n",
    "\n",
    "    tracks[movie][itrack] = pd.concat([tracks[movie][itrack].reset_index(drop=True), track_shape_metrics.reset_index(drop=True)], axis = 1 )\n",
    "\n",
    "    tracks[movie][itrack] = pd.concat([tracks[movie][itrack].reset_index(drop=True), med_centroids.reset_index(drop=True)], axis = 1 )\n",
    "\n",
    "    tracks[movie][itrack] = pd.concat([tracks[movie][itrack].reset_index(drop=True), shape_features.reset_index(drop=True)], axis = 1 )\n",
    "\n",
    "###Save updated tracks\n",
    "\n",
    "with open(local_analysis_output_folder/'tracks_shape.pkl', 'wb') as handle:\n",
    "    pickle.dump(tracks, handle, protocol=2)\n",
    "\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-NHBYQXCqE_X",
    "tags": []
   },
   "source": [
    "# Analyze Tracking Output\n",
    "Various analysis cells, most of which do not require masks or use cell shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_-yEy7AgNiMQ",
    "tags": []
   },
   "source": [
    "## Analysis setup (required for all other cells in this section)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "05KDw3nuqLWf",
    "tags": []
   },
   "source": [
    "### Load tracks, QC status, movie info, and get QC'd tracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XlF1MebDZSWu",
    "tags": []
   },
   "outputs": [],
   "source": [
    "fetch_analysis();\n",
    "with open(qc_params_path,'rb') as handle:\n",
    "    qc_params:Dict[str,Any] = pickle.load(handle)\n",
    "with open(qc_tracks_path, 'rb') as handle:\n",
    "    qc_tracks:Dict[int,Dict[int,pd.DataFrame]] = pickle.load(handle, encoding='latin1')\n",
    "with open(track_params_path,'rb') as handle:\n",
    "    track_params:Dict[str,Any] = pickle.load(handle,encoding='latin1');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VADjWnAvVfMo",
    "tags": []
   },
   "source": [
    "### Calibrate (Input) Pixelsize, Frametime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jFUBmSafZSW2",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Velocity Calibration\n",
    "#@markdown unit should be plural for readability reasons\n",
    "distance_unit:str = \"microns\" #@param{type:\"string\"}\n",
    "distance_per_pixel = 1.625 #@param{type:\"number\"}\n",
    "\n",
    "#@markdown unit should be singular for readability reasons\n",
    "time_unit:str = \"min\" #@param{type:\"string\"}\n",
    "time_per_frame = 5#@param{type:\"number\"}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Calibrate Stage Names, Groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown Extract stage names\n",
    "#@markdown If the .nd file is unavailable, you can manually map movie numbers to stage position names:\n",
    "do_override_map = False #@param{type:\"boolean\"}\n",
    "if do_override_map:\n",
    "    try:\n",
    "        stage_map = None #@param{type:\"raw\"}\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        override_map = None\n",
    "if not do_override_map or not override_map:   \n",
    "    #@markdown Movie numbers are associated with gradient images by their stage position name; stage positions can be extracted from a p.nd file from metamorph. Supported formatting: {experiment}\n",
    "    nd_location = \"gs://optotaxisbucket/movies/{experiment}/{experiment}/p.nd\" #@param{type:\"string\"}\n",
    "\n",
    "    #@markdown Whether to strip suffix when formatting for experiment - important if the .nd file is in images/cellmasks/nucmasks folders\n",
    "    strip_suffix = True #@param{type:\"boolean\"}\n",
    "    nd_location = Path(format_path(nd_location,experiment=experiment,strip_suffix=strip_suffix));\n",
    "\n",
    "    nd_local = gcp_transfer_folder/\"nd\"/experiment/(nd_location.name);\n",
    "\n",
    "    !gsutil cp \"{gs_str(nd_location)}\" \"{gs_str(nd_local)}\"\n",
    "    \n",
    "    stage_map:Dict[int,str] = {k:v for k,v in StageDict(nd_local).items() if k in qc_tracks.keys()}\n",
    "    \n",
    "    nd_data:Union[Dict[str,str],None] = parseND(nd_local)\n",
    "else:\n",
    "    stage_map:Dict[int,str] = override_map\n",
    "    nd_data:Union[Dict[str,str],None] = None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Create Stage Groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Stage groups: dict of {Group name: [(stage name, movie #)]}.\n",
    "group_dict:Dict[str,List[Tuple[str,int]]] = {}\n",
    "\n",
    "\n",
    "def get_group_movies(groupname:str):\n",
    "    return list(map(itemgetter(1),group_dict[groupname]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Put stage names as their own groups\n",
    "group_dict.update({name:[(name,num)] for num,name in stage_map.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Automatically create movie range groups\n",
    "#@markdown will create groups of the name format movie{num}, movie{start}-movie{end}, and \"all\"\n",
    "ms = range(min(qc_tracks.keys()),max(qc_tracks.keys())+1)\n",
    "group_dict.update({f\"movie{i}\":[(stage_map[i],i)] for i in ms})\n",
    "group_dict.update({f\"movie{i}-{j}\":[(stage_map[k],k) for k in range(i,j+1)] for i in ms[:-1] for j in range(i+1,ms[-1]+1)})\n",
    "group_dict.update({\"all\":[(v,k) for k,v in stage_map.items()]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Automatially group by stage basename (control1, control2, etc = control)\n",
    "group_dict.update(group_stage_basenames(stage_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(group_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6UbI4CUJNEUe",
    "tags": []
   },
   "source": [
    "### Calibrate Stage Names, Groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 126
    },
    "id": "xsk1p6K-ZSW8",
    "outputId": "8fde1834-fa16-4c13-c983-802679c90079",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown Stage groups: dict of {Group name: [(stage name, movie #)]}.\n",
    "group_dict:Dict[str,List[Tuple[str,int]]] = {}\n",
    "\n",
    "\n",
    "def get_group_movies(groupname:str):\n",
    "    return list(map(itemgetter(1),group_dict[groupname]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Put stage names as their own groups\n",
    "group_dict.update({name:[(name,num)] for num,name in stage_map.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@markdown Automatically create movie range groups\n",
    "#@markdown will create groups of the name format movie{num}, movie{start}-movie{end}, and \"all\"\n",
    "ms = range(min(qc_tracks.keys()),max(qc_tracks.keys())+1)\n",
    "group_dict.update({f\"movie{i}\":[(stage_map[i],i)] for i in ms})\n",
    "group_dict.update({f\"movie{i}-{j}\":[(stage_map[k],k) for k in range(i,j+1)] for i in ms[:-1] for j in range(i+1,ms[-1]+1)})\n",
    "group_dict.update({\"all\":[(v,k) for k,v in stage_map.items()]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['control1', 'control2', 'control3', 'control4', 'down1', 'down2', 'down3', 'down4', 'up1', 'up2', 'up3', 'up4']\n"
     ]
    }
   ],
   "source": [
    "#@markdown Automatially group by stage basename (control1, control2, etc = control)\n",
    "group_dict.update(group_stage_basenames(stage_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'control1': [('control1', 1)], 'control2': [('control2', 2)], 'control3': [('control3', 3)], 'control4': [('control4', 4)], 'down1': [('down1', 5)], 'down2': [('down2', 6)], 'down3': [('down3', 7)], 'down4': [('down4', 8)], 'up1': [('up1', 9)], 'up2': [('up2', 10)], 'up3': [('up3', 11)], 'up4': [('up4', 12)], 'movie1': [('control1', 1)], 'movie2': [('control2', 2)], 'movie3': [('control3', 3)], 'movie4': [('control4', 4)], 'movie5': [('down1', 5)], 'movie6': [('down2', 6)], 'movie7': [('down3', 7)], 'movie8': [('down4', 8)], 'movie9': [('up1', 9)], 'movie10': [('up2', 10)], 'movie11': [('up3', 11)], 'movie12': [('up4', 12)], 'movie1-2': [('control1', 1), ('control2', 2)], 'movie1-3': [('control1', 1), ('control2', 2), ('control3', 3)], 'movie1-4': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4)], 'movie1-5': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5)], 'movie1-6': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6)], 'movie1-7': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7)], 'movie1-8': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'movie1-9': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie1-10': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie1-11': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie1-12': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie2-3': [('control2', 2), ('control3', 3)], 'movie2-4': [('control2', 2), ('control3', 3), ('control4', 4)], 'movie2-5': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5)], 'movie2-6': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6)], 'movie2-7': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7)], 'movie2-8': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'movie2-9': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie2-10': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie2-11': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie2-12': [('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie3-4': [('control3', 3), ('control4', 4)], 'movie3-5': [('control3', 3), ('control4', 4), ('down1', 5)], 'movie3-6': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6)], 'movie3-7': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7)], 'movie3-8': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'movie3-9': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie3-10': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie3-11': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie3-12': [('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie4-5': [('control4', 4), ('down1', 5)], 'movie4-6': [('control4', 4), ('down1', 5), ('down2', 6)], 'movie4-7': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7)], 'movie4-8': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'movie4-9': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie4-10': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie4-11': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie4-12': [('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie5-6': [('down1', 5), ('down2', 6)], 'movie5-7': [('down1', 5), ('down2', 6), ('down3', 7)], 'movie5-8': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'movie5-9': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie5-10': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie5-11': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie5-12': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie6-7': [('down2', 6), ('down3', 7)], 'movie6-8': [('down2', 6), ('down3', 7), ('down4', 8)], 'movie6-9': [('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9)], 'movie6-10': [('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie6-11': [('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie6-12': [('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie7-8': [('down3', 7), ('down4', 8)], 'movie7-9': [('down3', 7), ('down4', 8), ('up1', 9)], 'movie7-10': [('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10)], 'movie7-11': [('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie7-12': [('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie8-9': [('down4', 8), ('up1', 9)], 'movie8-10': [('down4', 8), ('up1', 9), ('up2', 10)], 'movie8-11': [('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11)], 'movie8-12': [('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie9-10': [('up1', 9), ('up2', 10)], 'movie9-11': [('up1', 9), ('up2', 10), ('up3', 11)], 'movie9-12': [('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'movie10-11': [('up2', 10), ('up3', 11)], 'movie10-12': [('up2', 10), ('up3', 11), ('up4', 12)], 'movie11-12': [('up3', 11), ('up4', 12)], 'all': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4), ('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8), ('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)], 'control': [('control1', 1), ('control2', 2), ('control3', 3), ('control4', 4)], 'down': [('down1', 5), ('down2', 6), ('down3', 7), ('down4', 8)], 'up': [('up1', 9), ('up2', 10), ('up3', 11), ('up4', 12)]}\n"
     ]
    }
   ],
   "source": [
    "print(group_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6UbI4CUJNEUe",
    "tags": []
   },
   "source": [
    "## Export Tracks for Chemotaxis Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 126
    },
    "id": "xsk1p6K-ZSW8",
    "outputId": "8fde1834-fa16-4c13-c983-802679c90079",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Creates .txt file with tab delimited track info for import into the [Ibidi chemotaxis tool](https://ibidi.com/chemotaxis-analysis/171-chemotaxis-and-migration-tool.html)\n",
    "\n",
    "#@markdown ### The data will have distances of 1 of the unit selected in calibration, but the frame interval needs to be set in the tool itself\n",
    "\n",
    "\n",
    "def export_trackdata(selected_movies,filename):\n",
    "    print(\"processing movies\",selected_movies)\n",
    "\n",
    "    centertype = track_params['centroidtype'];\n",
    "    centerx = centertype + 'x';\n",
    "    centery = centertype + 'y';\n",
    "\n",
    "    combined_tracks:List[pd.DataFrame] = [];\n",
    "    trackmap = [];\n",
    "    for mov in selected_movies:\n",
    "        combined_tracks.extend(qc_tracks[mov].values());\n",
    "        trackmap.extend([(mov,k) for k in qc_tracks[mov].keys()]);\n",
    "    trackids = (dict([(x+1,y) for x,y in enumerate(trackmap)]))\n",
    "\n",
    "\n",
    "    trackdata = [];\n",
    "    m = master_bar(combined_tracks)\n",
    "    for id,tr in enumerate(m):\n",
    "        for _,row in progress_bar(tr.iterrows(),parent=m,total=tr.shape[0]):\n",
    "            slic = int(row['frame']);\n",
    "            x = int(row[centerx]*distance_per_pixel);\n",
    "            y = int(row[centery]*distance_per_pixel);\n",
    "            trackdata.append((id+1,slic,x,y));\n",
    "    with open(filename,'w') as f:\n",
    "        f.write(\"Combined track data from movies: \" + \", \".join([str(m) for m in selected_movies]) + \", track ids:\" + str(trackids) + \"\\n\");\n",
    "        f.writelines([f\"{i+1}\\t{tid}\\t{slid}\\t{x}\\t{y}\\n\" for i,(tid,slid,x,y) in enumerate(progress_bar(trackdata))]);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown save trackdata file for group(s):\n",
    "\n",
    "\n",
    "#@markdown groups to use in experiment; quick guide: single movie = \"movie{num}\", movie range = \"movie{start}-{end}\", all: \"all\"\n",
    "groups = [\"control\",\"up\",\"down\"] #@param\n",
    "\n",
    "#@markdown ### Saving settings\n",
    "#@markdown By default will use the format trackdata_{groupname}_{no_cell_collection/using_cell_collection}.txt\n",
    "\n",
    "#@markdown whether to override this default behavior and use the specified filename instead: (will be saved in segmentation analysis regardless)\n",
    "#@markdown \"{collection}\", \"{groupname}\", \"{smoothing}\" will all be available for formatting if custom\n",
    "\n",
    "do_custom_filename = True #@param{type:\"boolean\"}\n",
    "custom_filename = \"trackdata_{groupname}_{smoothing}.txt\" #@param{type:\"string\"}\n",
    "\n",
    "filename = custom_filename if do_custom_filename else \"trackdata_{groupname}.txt\";\n",
    "\n",
    "for group in groups:\n",
    "    selected_movies = get_group_movies(group)\n",
    "    fname = format_path(filename,collection=track_params['do_cell_collection'],groupname=group,qc_params=qc_params);\n",
    "    fname = local_analysis_output_folder/fname;\n",
    "    export_trackdata(selected_movies,fname)\n",
    "\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZTRxRC_NrjGh",
    "tags": []
   },
   "source": [
    "## Calculate FMIs, Persistence, Average Velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 37
    },
    "id": "uK6rolqRZSW-",
    "outputId": "176216cc-8983-4e2f-eafd-00bd82d21d54",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Calculate FMIs, Velocity, and Displacement for each track (code hidden)\n",
    "\n",
    "def dist(x1,x2,y1,y2):\n",
    "  return math.sqrt((x1-x2)**2+(y1-y2)**2)\n",
    "\n",
    "##FMI: two axis (x and y) - total displacement in the axis across a track divided by the total distance traveled by the cell\n",
    "FMI: Dict[int,Dict[int,Tuple[float,float]]] = {}; #{movie, {trackid:(FMI.x,FMI.y)}}\n",
    "Persistence: Dict[int,Dict[int,float]] = {}; #{movie, {trackid:Persistence}}\n",
    "trackVelocity: Dict[int,Dict[int,Tuple[float,float,float]]] = {}; #{movie, {trackid:(velocityX,velocityY,velocityMag)}}\n",
    "trackDisplacement: Dict[int,Dict[int,Tuple[float,float,float]]] = {}; #{movie, {trackid:(displacementX,displacementY,displacementDist)}}\n",
    "trackLength: Dict[int,Dict[int,float]] = {}; #{movie, {trackid:tracklength}}\n",
    "trackTime: Dict[int,Dict[int,int]] = {}; #{movie, {trackid:trackTime}}\n",
    "\n",
    "avgFMI: Dict[int,Tuple[float,float]] = {}; #{movie,(avgX,avgY)};\n",
    "avgPersistence: Dict[int,float] = {}; #{movie,average};\n",
    "avgVelocity: Dict[int,Tuple[float,float,float]] = {}; #{movie,(averageX,averagyY,averageMag)};\n",
    "avgDisplacement: Dict[int,Tuple[float,float,float]] = {}; #{movie,(avgX,avgY,avgDist)}\n",
    "avgTracklength: Dict[int,float] = {}; #{movie,average}\n",
    "avgTracktime: Dict[int,float] = {}; #{movie,average}\n",
    "\n",
    "centertype = track_params['centroidtype'];\n",
    "centerx = centertype + 'x';\n",
    "centery = centertype + 'y';\n",
    "\n",
    "bar = master_bar(qc_tracks.items())\n",
    "for movie,tracks in bar:\n",
    "    FMI[movie] = {};\n",
    "    Persistence[movie] = {};\n",
    "    trackVelocity[movie] = {};\n",
    "    trackDisplacement[movie] = {}\n",
    "    trackLength[movie] = {}\n",
    "    trackTime[movie] = {}\n",
    "\n",
    "    numPoints = 0;\n",
    "    FMI_accX = 0;\n",
    "    FMI_accY = 0;\n",
    "    Persistence_acc = 0;\n",
    "    Velocity_acc = np.array([0,0,0]);\n",
    "    Displacement_acc = np.array([0,0,0]);\n",
    "    Tracklength_acc = 0;\n",
    "    Tracktime_acc = 0\n",
    "\n",
    "    for id,data in progress_bar(tracks.items(),parent=bar):\n",
    "        numPoints += 1;\n",
    "\n",
    "        start = data.iloc[0];\n",
    "        end = data.iloc[-1];\n",
    "    \n",
    "        ##Get accumulated distance (total movement within track)\n",
    "        accDist = 0;\n",
    "        accVel = np.array([0,0,0]);\n",
    "        prevpos = (start[centerx]*distance_per_pixel,start[centery]*distance_per_pixel);\n",
    "        for x,y in zip(data.iloc[1:][centerx],data.iloc[1:][centery]):\n",
    "            x *= distance_per_pixel;\n",
    "            y *= distance_per_pixel;\n",
    "            x1,y1 = prevpos;\n",
    "            accVel = np.sum([accVel,[x-x1,y-y1,dist(x,x1,y,y1)]],axis=0);\n",
    "            accDist += math.sqrt((y-y1)**2+(x-x1)**2);\n",
    "            prevpos = (x,y);\n",
    "    \n",
    "        ##Get vertical, horizontal displacement\n",
    "        xDisp = (end[centerx] - start[centerx])*distance_per_pixel;\n",
    "        yDisp = (end[centery] - start[centery])*distance_per_pixel;\n",
    "\n",
    "        ##Get individual cell FMI\n",
    "        xMI = xDisp/accDist;\n",
    "        yMI = yDisp/accDist;\n",
    "\n",
    "        FMI[movie][id] = (xMI,yMI);\n",
    "\n",
    "        ## Get net cell distance\n",
    "        netDist = math.sqrt(xDisp**2 + yDisp**2);\n",
    "\n",
    "        ## Get Persistence\n",
    "        direct = netDist/accDist;\n",
    "        Persistence[movie][id] = direct;\n",
    "\n",
    "        ##Get Average Velocity\n",
    "        avgTrackVel = accVel/(len(data)*time_per_frame);\n",
    "        trackVelocity[movie][id] = tuple(avgTrackVel);\n",
    "        \n",
    "        #Get Displacements\n",
    "        trackDisplacement[movie][id] = (xDisp,yDisp,netDist)\n",
    "        \n",
    "        #Get Tracklength\n",
    "        trackLength[movie][id] = accDist\n",
    "        \n",
    "        #Get Tracktime\n",
    "        trackTime[movie][id] = len(data)*time_per_frame\n",
    "\n",
    "        ##Accumulate FMI, Persistence, and Velocity\n",
    "        FMI_accX += xMI;\n",
    "        FMI_accY += yMI;\n",
    "        Persistence_acc += direct;\n",
    "        Velocity_acc = np.sum([Velocity_acc,avgTrackVel],axis=0)\n",
    "        Displacement_acc = np.sum([Displacement_acc,(xDisp,yDisp,netDist)],axis=0)\n",
    "        Tracklength_acc += accDist\n",
    "        Tracktime_acc += len(data)*time_per_frame\n",
    "\n",
    "    if (numPoints > 0):\n",
    "        avgFMI[movie] = (FMI_accX/numPoints,FMI_accY/numPoints);\n",
    "        avgPersistence[movie] = Persistence_acc/numPoints;\n",
    "        avgVelocity[movie] = tuple(Velocity_acc/numPoints);\n",
    "        avgDisplacement[movie] = tuple(Displacement_acc/numPoints);\n",
    "        avgTracklength[movie] = Tracklength_acc/numPoints\n",
    "        avgTracktime[movie] = Tracktime_acc/numPoints\n",
    "    else:\n",
    "        avgFMI[movie] = (0,0);\n",
    "        avgPersistence[movie] = 0;\n",
    "        avgVelocity[movie] = (0,0,0);\n",
    "        avgDisplacement[movie] = (0,0,0);\n",
    "        avgTracklength[movie] = 0\n",
    "        avgTracktime[movie] = 0\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "OZ1HVWeTZSXB",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Save Track Analysis to csv\n",
    "\n",
    "#@markdown ### Output path in the GCP bucket (csv file)\n",
    "gcp_out_path = \"track_analysis_{smoothing}.csv\" #@param {type:\"string\"}\n",
    "gcp_out_path = Path(format_path(gcp_out_path,collection=track_params['do_cell_collection'],qc_params=qc_params));\n",
    "if not str(gcp_out_path).endswith(\".csv\"):\n",
    "    gcp_out_path = gcp_out_path.with_suffix('.csv');\n",
    "\n",
    "#@markdown ##### Treat the above path as relative to the experiment's analysis output folder?\n",
    "in_analysis = True #@param {type:\"boolean\"}\n",
    "\n",
    "local_path = local_analysis_output_folder / gcp_out_path if in_analysis else temp_folder / Path(gcp_out_path).name\n",
    "gcp_out_path = gcp_analysis_output_folder / gcp_out_path if in_analysis else Path(gcp_out_path);\n",
    "\n",
    "with open(local_path,\"w\", newline='') as file:\n",
    "  fieldnames = ['movie', 'trackid','FMI.x','FMI.y',f'Velocity.x ({distance_unit}/{time_unit})',f'Velocity.y ({distance_unit}/{time_unit})',f'Speed ({distance_unit}/{time_unit})','Persistence',f'Displacement.x ({distance_unit})',f'Displacement.y ({distance_unit})',f'Displacement Distance ({distance_unit})',f'Tracklength ({distance_unit})',f'Track time ({time_unit})'];\n",
    "  writer = csv.DictWriter(file, fieldnames=fieldnames);\n",
    "  writer.writeheader()\n",
    "  for movie in qc_tracks:\n",
    "    writer.writerow(dict(zip(fieldnames,\n",
    "                             [movie,\"average\",\n",
    "                              avgFMI[movie][0],\n",
    "                              avgFMI[movie][1],\n",
    "                              avgVelocity[movie][0],\n",
    "                              avgVelocity[movie][1],\n",
    "                              avgVelocity[movie][2],\n",
    "                              avgPersistence[movie],\n",
    "                              avgDisplacement[movie][0],\n",
    "                              avgDisplacement[movie][1],\n",
    "                              avgDisplacement[movie][2],\n",
    "                              avgTracklength[movie],\n",
    "                              avgTracktime[movie]\n",
    "                             ])));\n",
    "  for movie,tracks in progress_bar(qc_tracks.items()):\n",
    "    for id in tracks:\n",
    "      writer.writerow(dict(zip(fieldnames,\n",
    "                             [movie,id,\n",
    "                              FMI[movie][id][0],\n",
    "                              FMI[movie][id][1],\n",
    "                              trackVelocity[movie][id][0],\n",
    "                              trackVelocity[movie][id][1],\n",
    "                              trackVelocity[movie][id][2],\n",
    "                              Persistence[movie][id],\n",
    "                              trackDisplacement[movie][id][0],\n",
    "                              trackDisplacement[movie][id][1],\n",
    "                              trackDisplacement[movie][id][2],\n",
    "                              trackLength[movie][id],\n",
    "                              trackTime[movie][id],\n",
    "                             ])));\n",
    "print(\"file written, copying to bucket...\")\n",
    "!gsutil -m cp \"{gs_str(local_path)}\" \"{gs_str(gcp_out_path)}\"\n",
    "print(\"file copied\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fws42Mb--YQT",
    "tags": []
   },
   "source": [
    "## Plot Tracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 191
    },
    "id": "t8xvwuvxZSXE",
    "outputId": "86f0f823-fe9e-46df-b385-e3934cb881b8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Compute track points\n",
    "\n",
    "#@markdown whether to shift all the tracks so their centers are at the origin, or to plot their positions how they were in 2d space\n",
    "from_origin = False #@param {type:\"boolean\"}\n",
    "\n",
    "centertype = track_params['centroidtype'];\n",
    "centerx = centertype + 'x';\n",
    "centery = centertype + 'y';\n",
    "\n",
    "origin_adjusted_tracks:Dict[int,Dict[int,List[Tuple[float,float]]]] = {}; #dict of {movie: dict of {trackid:[list of (x,y)]}}\n",
    "\n",
    "bar = master_bar(qc_tracks.items())\n",
    "for movie,tracks in bar:\n",
    "  print(movie);\n",
    "  origin_adjusted_tracks[movie] = {};\n",
    "\n",
    "  for id,data in progress_bar(tracks.items(),parent=bar):\n",
    "    if from_origin:\n",
    "      start = data.iloc[0];\n",
    "    else:\n",
    "      start = {centerx:0,centery:0};\n",
    "    points = [];\n",
    "    for x,y in zip(data.iloc[1:][centerx],data.iloc[1:][centery]):\n",
    "      points.append(((x-start[centerx]),(y-start[centery])));\n",
    "    origin_adjusted_tracks[movie][id]=points;\n",
    "\n",
    "\n",
    "# print(len(origin_adjusted_tracks[5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 470
    },
    "id": "vLpIU5oFZSXF",
    "outputId": "3724bda8-95e2-40d1-93c0-f8f2c95433cd",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Plot movie tracks\n",
    "#@markdown Movie group to plot (if \"None\" just prepares the tracking function)\n",
    "movie_group = \"None\" #@param{type:\"str\"}\n",
    "\n",
    "#@markdown What to title the plot; can be formatted with {groupname}, {experiment}, {collection}, {smoothing}\n",
    "plot_title = \"{groupname} {smoothing} {experiment_nostrip} tracks\"\n",
    "\n",
    "#@markdown whether to plot each track in different colors (if false, all will be black);\n",
    "do_color = True #@param{type:\"boolean\"}\n",
    "\n",
    "#@markdown whether to show a legend mapping track id to color\n",
    "do_legend = True #@param{type:\"boolean\"}\n",
    "\n",
    "#@markdown whether to put a dot at the beginning of each track\n",
    "plot_start = True\n",
    "\n",
    "#@markdown set plot boundaries for uniform scaling; if None will use matplotlib's autoscaling\n",
    "#@markdown format: [xmin,xmax,ymin,ymax]\n",
    "#@markdown note: input bounds in picture coordinates; trackplot y axis is inverted, so y tick numbers are incorrect\n",
    "plot_bounds = [-50,1300,-50,1100]\n",
    "\n",
    "\n",
    "def plot_movies(group=movie_group,legend=do_legend,color=do_color,startdot=plot_start):\n",
    "    print(\"Plotting group:\",group)\n",
    "    sel_movies = get_group_movies(group);\n",
    "    figure = plt.figure()\n",
    "    ax = plt.axes()\n",
    "    xs = [];\n",
    "    ys = [];\n",
    "    labels = [];\n",
    "    for movie in sel_movies:\n",
    "        movie_tracks = origin_adjusted_tracks[movie];\n",
    "        xs += [[t[0] for t in track] for track in movie_tracks.values()];\n",
    "        ys += [[-t[1] for t in track] for track in movie_tracks.values()];\n",
    "        if len(sel_movies) == 1:\n",
    "            labels += [tnum for tnum in movie_tracks.keys()];\n",
    "        else:\n",
    "            labels += [f'm{movie}_t{tnum}' for tnum in movie_tracks.keys()];\n",
    "\n",
    "    colors = cm.rainbow(np.linspace(0, 1, len(ys)))\n",
    "    for (x,y,n,c) in zip(xs,ys,labels,colors):\n",
    "        if do_color:\n",
    "            ax.plot(x, y, color=c,linewidth=1,label=n);\n",
    "        else:\n",
    "            ax.plot(x,y,color = 'black', linewidth=1,label=n);\n",
    "    if startdot:\n",
    "        for (x,y) in zip(xs,ys):\n",
    "            ax.plot(x[0],y[0], color='black',marker='.')\n",
    "\n",
    "    ax.set_xlabel(f\"X distance traveled ({distance_unit})\")\n",
    "    ax.set_ylabel(f\"Y distance traveled ({distance_unit})\")\n",
    "    ax.set_title(format_path(plot_title,experiment=experiment,collection=track_params[\"do_cell_collection\"],groupname=group,qc_params=qc_params))\n",
    "    ax.plot(0,0,color= 'black' if do_color else 'red',marker='o',markersize=3);\n",
    "    ax.axis('equal')\n",
    "    \n",
    "    if plot_bounds is not None and len(plot_bounds) == 4:\n",
    "        ax.set_aspect('auto')\n",
    "        xb = (plot_bounds[0],plot_bounds[1])\n",
    "        yb = (-plot_bounds[3],-plot_bounds[2])\n",
    "        ax.set_xlim(xmin=xb[0],xmax=xb[1])\n",
    "        ax.set_ylim(ymin=yb[0],ymax=yb[1])\n",
    "        # print(\"Plot xbounds:\",xb,\"ybounds:\",yb)\n",
    "        \n",
    "\n",
    "    if do_legend:\n",
    "        figure.legend(loc=\"center left\");\n",
    "\n",
    "    return figure\n",
    "if movie_group != \"None\":\n",
    "    _ = plot_movies()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2a41jzoVZSXG",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Save plot(s) for group(s) of movies\n",
    "\n",
    "#@markdown Groups to save: (list of groupname strings)\n",
    "movie_groups:List[str] = [\"control\",\"up\",\"down\",*[stage_map[i] for i in qc_tracks]] #@param\n",
    "\n",
    "#@markdown By default will use the format trackplot_{groupname}.png\n",
    "\n",
    "#@markdown whether to override this default behavior and use the specified filename instead: (will be saved in segmentation analysis regardless)\n",
    "#@markdown \"{collection}\", \"{groupname}\", \"{smoothing}\" will all be available for formatting if custom\n",
    "do_custom_filename = True #@param{type:\"boolean\"}\n",
    "custom_filename = \"trackplot_{groupname}_{smoothing}.png\" #@param{type:\"string\"}\n",
    "\n",
    "filename = custom_filename if do_custom_filename else \"trackplot_{groupname}.png\";\n",
    "for group in progress_bar(movie_groups):\n",
    "    # group_movies = get_group_movies(group)\n",
    "    fig = plot_movies(group)\n",
    "    \n",
    "    fname = format_path(filename,collection=track_params['do_cell_collection'],groupname=group,qc_params=qc_params);\n",
    "    fig.savefig(local_analysis_output_folder/fname,bbox_inches='tight');\n",
    "    # fig.show()\n",
    "\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s_wsyfSVjJpK",
    "tags": []
   },
   "source": [
    "## Track Gradient Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pzIPbgGiZSXH",
    "outputId": "b06ac583-3de6-4920-f6ea-3d9884a495ff",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown # Input paths to gradients (and fetch from gcp)\n",
    "\n",
    "#@markdown Whether to trim suffixes (\" $[suffix]\") from the experiment name); if no suffixes, this will do nothing\n",
    "strip_suffix = True #@param{\"type\":\"boolean\"}\n",
    "\n",
    "#@markdown Folder in gcp where gradients are stored. \"{experiment}\" will be formatted with the experiment name. Include the gs://\n",
    "gradient_folder = \"gs://optotaxisbucket/movies/{experiment}/Gradient Calibration\" #@param{type:\"string\"}\n",
    "gradient_folder = Path(format_path(gradient_folder,experiment=experiment,strip_suffix=strip_suffix));\n",
    "\n",
    "#@markdown Supported formatting: {experiment}, {movie} (number), or {stage} (from .nd file or override map) for each position used\n",
    "gradient_filename = \"{stage}.tif\" #@param{type:\"string\"}\n",
    "gradient_filename = format_path(gradient_filename,experiment=experiment,strip_suffix=strip_suffix);\n",
    "\n",
    "gradient_local = gcp_transfer_folder/\"gradients\"/experiment\n",
    "if not gradient_local.exists():\n",
    "  os.makedirs(gradient_local);\n",
    "\n",
    "for movie,name in stage_map.items():\n",
    "  gf = format_path(gradient_folder/gradient_filename,stage=name,movie=movie)\n",
    "  print(gf)\n",
    "  with capture_output():\n",
    "    !gsutil cp \"{gs_str(gf)}\" \"{gs_str(gradient_local)}/\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EoVOi86FrTau",
    "tags": []
   },
   "source": [
    "### Plot Tracks over gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Export Persistence, Alignment, Velocity with Gradient Steepness, Brightness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_gradient(6).show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "9GKMcq9nZSXO",
    "outputId": "69953e2a-82c1-4539-8c9a-5e78686c5c8c",
    "tags": []
   },
   "source": [
    "#@markdown ## Calculate Frame-by-Frame Metrics with Gradient Steepness X,Y; Intensity\n",
    "\n",
    "#@markdown ### Flattened save location (in analysis): Combined csv file with every frame from every movie\n",
    "#@markdown #### Supported formatting: {collection}, {experiment}, {qc_params}\n",
    "flattened_save_location = \"gradient_track_flattened_analysis.csv\" #@param {\"type\":\"string\"}\n",
    "flattened_save_path = Path(format_path(flattened_save_location,collection=track_params['do_cell_collection'],experiment=experiment,qc_params=qc_params))\n",
    "if flattened_save_path.suffix != \".csv\":\n",
    "  if flattened_save_path.suffix != \"\":\n",
    "    print(\"warning: flattened analysis must be .csv file; changing extension to match\");\n",
    "  flattened_save_path.suffix = \".csv\";\n",
    "\n",
    "\n",
    "\n",
    "#@markdown ### Shaped (python dicts) data save location (in analysis): .pkl file with trackdata organized by movie and track id\n",
    "#@markdown #### Supported formatting: {collection}, {experiment}, {smoothing}\n",
    "shaped_save_location = \"gradient_track_shaped_analysis.pkl\" #@param {\"type\":\"string\"}\n",
    "shaped_save_path = Path(format_path(shaped_save_location,collection=track_params['do_cell_collection'],experiment=experiment,qc_params=qc_params))\n",
    "if shaped_save_path.suffix != \".pkl\":\n",
    "  if shaped_save_path.suffix != \"\":\n",
    "    print(\"warning: shaped analysis must be .pkl file; changing extension to match\");\n",
    "  shaped_save_path.suffix = \".pkl\";\n",
    "\n",
    "\n",
    "\n",
    "movies = qc_tracks.keys()\n",
    "\n",
    "centertype = track_params['centroidtype'];\n",
    "centerx = centertype + 'x';\n",
    "centery = centertype + 'y';\n",
    "\n",
    "def pairwise(iterable):\n",
    "    \"s -> (s0,s1), (s1,s2), (s2, s3), ...\"\n",
    "    a, b = itertools.tee(iterable)\n",
    "    next(b, None)\n",
    "    return zip(a, b)\n",
    "\n",
    "def triplewise(iterable):\n",
    "    \"Return overlapping triplets from an iterable\"\n",
    "    # triplewise('ABCDEFG') --> ABC BCD CDE DEF EFG\n",
    "    for (a, _), (b, c) in pairwise(pairwise(iterable)):\n",
    "        yield a, b, c\n",
    "\n",
    "\n",
    "def dist(x1,x2,y1,y2):\n",
    "  return math.sqrt((x1-x2)**2 + (y1-y2)**2)\n",
    "\n",
    "def mag(x,y):\n",
    "  return math.sqrt(x**2+y**2);\n",
    "\n",
    "shaped_data:Dict[int,Dict[int,pd.DataFrame]] = {}\n",
    "flattened_data:pd.DataFrame = pd.DataFrame(columns=[\"movie\",\"frame\",\"trackid\",\"label\",centerx,centery,\"gradient.x\",\"gradient.y\",\"gradient intensity\",\"velocity.x\",\"velocity.y\",\"speed\",\"persistence\",\"dTheta\"])\n",
    "flat_datalist = [];\n",
    "parent = master_bar(movies)\n",
    "for movie in parent:\n",
    "  #get gradient information\n",
    "  gradient = imread(format_path(gradient_local/gradient_filename,stage=stage_map[movie]));\n",
    "  gradient = cv2.GaussianBlur(gradient,(15,15),0);\n",
    "\n",
    "  xAvg = np.average(gradient,0);\n",
    "  xX = np.arange(gradient.shape[1])\n",
    "  xInterpolation = scipy.interpolate.CubicSpline(xX,xAvg)\n",
    "\n",
    "  yAvg = np.average(gradient,1);\n",
    "  yX = np.arange(gradient.shape[0])\n",
    "  yInterpolation = scipy.interpolate.CubicSpline(yX,yAvg)\n",
    "  \n",
    "  def grad_xgrad(x,y):\n",
    "    return xInterpolation(x,1);\n",
    "\n",
    "  def grad_ygrad(x,y):\n",
    "    return yInterpolation(y,1);\n",
    "\n",
    "  def grad_mid(x,y):\n",
    "    return yInterpolation(y); #THIS IS ASSUMING CONSTANT / SMALL X DEVIATION, ONLY WORKS FOR VERTICAL GRADIENTS\n",
    "\n",
    "\n",
    "  shaped_data[movie] = {};\n",
    "  \n",
    "  ### velocity is from previous frame to current, persistence + dTheta is dP of previous frame and dP of next\n",
    "  for tid, track in progress_bar(qc_tracks[movie].items(),parent=parent):\n",
    "    trackData = [];    \n",
    "    tracklength = track.shape[0]\n",
    "    # print(tracklength);\n",
    "\n",
    "    #1st-order (gradient)\n",
    "    for x,y,(_,frame) in zip(track[centerx],track[centery],track.iterrows()):\n",
    "      frame = frame.copy();\n",
    "      frame['gradient.x'] = grad_xgrad(x,y);\n",
    "      frame['gradient.y'] = grad_ygrad(x,y);\n",
    "      frame['gradient intensity'] = grad_mid(x,y);\n",
    "      trackData.append(frame)\n",
    "    \n",
    "    #second-order (velocity)\n",
    "    for frame1,frame2 in pairwise(trackData):\n",
    "      x1,y1 = frame1[centerx],frame1[centery];\n",
    "      x2,y2 = frame2[centerx],frame2[centery];\n",
    "\n",
    "      frame2['velocity.x'] = (x2-x1)*distance_per_pixel/time_per_frame;\n",
    "      frame2['velocity.y'] = (y2-y1)*distance_per_pixel/time_per_frame;\n",
    "      frame2['speed'] = dist(x1,x2,y1,y2)*distance_per_pixel/time_per_frame;\n",
    "\n",
    "    #third-order (persistence, dTheta)\n",
    "    for frame1,frame2,frame3 in triplewise(trackData):\n",
    "      x1,y1 = frame1[centerx],frame1[centery];\n",
    "      x2,y2 = frame2[centerx],frame2[centery];  \n",
    "      x3,y3 = frame3[centerx],frame3[centery];\n",
    "      \n",
    "      v1 = (x2-x1,y2-y1);\n",
    "      v2 = (x3-x2,y3-y2);\n",
    "      v3 = (x3-x1,y3-y1);\n",
    "\n",
    "      disp = mag(*v3);\n",
    "      l1 = mag(*v1);\n",
    "      l2 = mag(*v2);\n",
    "    \n",
    "\n",
    "      length = l1 + l2;\n",
    "      persistence = disp/length if length > 0.01 else float('nan');\n",
    "\n",
    "      frame2['persistence'] = persistence;\n",
    "\n",
    "\n",
    "    trackframe = pd.DataFrame(trackData);\n",
    "    flat_datalist.append(trackframe);\n",
    "    shaped_data[movie][tid] = trackframe\n",
    "\n",
    "flattened_data = pd.concat(flat_datalist);\n",
    "# print(flattened_data);\n",
    "\n",
    "out = local_analysis_output_folder/flattened_save_path;\n",
    "out.parent.mkdir(parents=True,exist_ok=True);\n",
    "flattened_data.to_csv(out);\n",
    "\n",
    "out = local_analysis_output_folder/shaped_save_path;\n",
    "out.parent.mkdir(parents=True,exist_ok=True);\n",
    "with open(out,'wb') as f:\n",
    "  pickle.dump(out,f);\n",
    "\n",
    "push_analysis();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Plot Tracks over gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "vX44x3P5ZSXK"
   },
   "outputs": [],
   "source": [
    "#@markdown # Compute Track Points and Gradient \n",
    "\n",
    "centertype = track_params['centroidtype'];\n",
    "centerx = centertype + 'x';\n",
    "centery = centertype + 'y';\n",
    "\n",
    "\n",
    "def plot_gradient(movie:int,do_color=True,do_startmarker=True):\n",
    "  plt.figure()\n",
    "  gradient = imread(format_path(gradient_local/gradient_filename,stage=stage_map[movie]));\n",
    "  plt.imshow(gradient);\n",
    "  \n",
    "  trackPoints = {};\n",
    "  for id,data in progress_bar(qc_tracks[movie].items()):\n",
    "    trackPoints[id] = list(zip(data[centerx],data[centery]));\n",
    "  \n",
    "  xs = [[t[0] for t in track] for track in trackPoints.values()];\n",
    "  ys = [[t[1] for t in track] for track in trackPoints.values()];\n",
    "  labels = [tnum for tnum in trackPoints.keys()];\n",
    "  \n",
    "  colors = cm.rainbow(np.linspace(0, 1, len(ys)))\n",
    "  for (x,y,n,c) in zip(xs,ys,labels,colors):\n",
    "\n",
    "    if do_color:\n",
    "      plt.plot(x, y, color=c,linewidth=1,label=n);\n",
    "    else:\n",
    "      plt.plot(x,y,color = 'black', linewidth=1,label=n);\n",
    "\n",
    "    if do_startmarker:\n",
    "      plt.plot(x[0],y[0],marker='o',color='black',markersize=2);\n",
    "\n",
    "  plt.xlabel(f\"X distance traveled ({distance_unit})\")\n",
    "  plt.ylabel(f\"Y distance traveled ({distance_unit})\")\n",
    "  # plt.plot(0,0,color= 'black' if do_color else 'red',marker='o',markersize=3);\n",
    "  plt.axis('equal')\n",
    "  \n",
    "  fig = plt.gcf();\n",
    "  return fig;\n",
    "      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "cgqWJvuQE8j3",
    "outputId": "c01fcef5-28cc-4734-d7eb-b6d5a5d75d6f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@markdown ## Save plots\n",
    "\n",
    "#@markdown By default will use the format gradient_trackplot_{movie}.png\n",
    "\n",
    "#@markdown Which movies to save; set as None to do them all.\n",
    "savemovies = [] #@param\n",
    "if savemovies is None:\n",
    "    savemovies = stage_map.keys();\n",
    "    \n",
    "\n",
    "#@markdown whether to override this default behavior and use the specified filename instead: (will be saved in segmentation analysis regardless)\n",
    "#@markdown \"{movie}\",\"{stagename}\",\"{smoothing}\" will be available for formatting if custom\n",
    "do_custom_filename = True #@param{type:\"boolean\"}\n",
    "custom_filename = \"gradient_trackplot_{stagename}.png\" #@param{type:\"string\"}\n",
    "\n",
    "for movie in savemovies:\n",
    "    filename = custom_filename if do_custom_filename else \"gradient_trackplot_{movie}.png\";\n",
    "    filename = format_path(filename,collection=track_params['do_cell_collection'],movie=movie,stagename=stage_map.get(movie,None),qc_params=qc_params);\n",
    "\n",
    "    try:\n",
    "        fig = plot_gradient(movie)\n",
    "        fig.savefig(local_analysis_output_folder/filename);\n",
    "    except IndexError:\n",
    "        print(\"movie\",movie,\"missing from tracks, skipping\")\n",
    "\n",
    "push_analysis();"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "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.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "1fab5b2657ca90398e2f7d6e349d246f9168433ae93312ee9b22cab62345ecaa"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}