{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from cellpose import models\n",
    "from cellpose.io import imread\n",
    "import glob\n",
    "from pathlib import Path\n",
    "from PIL import Image, ImageSequence\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import os.path\n",
    "from livecellx import core\n",
    "from livecellx.core import datasets\n",
    "from livecellx.core.datasets import LiveCellImageDataset, SingleImageDataset\n",
    "from skimage import measure\n",
    "from livecellx.core import SingleCellTrajectory, SingleCellStatic\n",
    "import numpy as np\n",
    "import os, json, cv2, random\n",
    "import cv2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SingleCellStatic: saving & loading "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "io_out_dir = Path(\"test_io_output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir_path = Path(\n",
    "    \"../datasets/test_data_STAV-A549/DIC_data\"\n",
    ")\n",
    "\n",
    "mask_dataset_path = Path(\"../datasets/test_data_STAV-A549/mask_data\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 png img file paths loaded;\n"
     ]
    }
   ],
   "source": [
    "mask_dataset = LiveCellImageDataset(mask_dataset_path, ext=\"png\")\n",
    "\n",
    "time2url = sorted(glob.glob(str((Path(dataset_dir_path) / Path(\"*_DIC.tif\")))))\n",
    "time2url = {i: path for i, path in enumerate(time2url)}\n",
    "dic_dataset = LiveCellImageDataset(time2url=time2url, ext=\"tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:06<00:00,  2.25s/it]\n"
     ]
    }
   ],
   "source": [
    "from skimage.measure import regionprops\n",
    "from livecellx.core.io_sc import prep_scs_from_mask_dataset\n",
    "single_cells = prep_scs_from_mask_dataset(mask_dataset, dic_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'label_in_mask': 2}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "single_cells[1].meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Writing single cells to json  \n",
    "You can check the documentation of `write_single_cells_json` for more details.  \n",
    "`return_list` controls whether or not to return json object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "SingleCellStatic.write_single_cells_json(single_cells, path=io_out_dir/\"single_cells_0.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc_json_list = SingleCellStatic.write_single_cells_json(single_cells, io_out_dir/\"single_cells_1.json\", dataset_dir=io_out_dir/\"dataset\", return_list=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "SingleCellStatic.write_single_cells_json(single_cells, io_out_dir/\"single_cells_2.json\", dataset_dir=io_out_dir/\"dataset\", return_list=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dataset = single_cells[0].img_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "constructing single cells from json dict: 100%|██████████| 42/42 [00:00<00:00, 5673.27it/s]\n"
     ]
    }
   ],
   "source": [
    "loaded_scs = SingleCellStatic.load_single_cells_json(io_out_dir/\"single_cells_0.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SingleCellStatic(id=fc9d1413-2399-4b6f-bf91-13ba4e05b34c, timeframe=0, bbox=[ 143.  978.  207. 1044.])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loaded_scs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fc9d1413-2399-4b6f-bf91-13ba4e05b34c\n",
      "402fbdcf-c641-45ca-a580-54fe5b7519e8\n"
     ]
    }
   ],
   "source": [
    "for sc in single_cells[:2]:\n",
    "    print(sc.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fc9d1413-2399-4b6f-bf91-13ba4e05b34c\n",
      "402fbdcf-c641-45ca-a580-54fe5b7519e8\n"
     ]
    }
   ],
   "source": [
    "for sc in loaded_scs[:2]:\n",
    "    print(sc.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sc in single_cells:\n",
    "    for loaded_sc in loaded_scs:\n",
    "        if sc.id == loaded_sc.id:\n",
    "            # compare contour, contours are np.array\n",
    "            assert np.allclose(sc.contour, loaded_sc.contour), f\"the difference between sc.contour and loaded_sc.contour is {sc.contour - loaded_sc.contour}, ids are {sc.id} and {loaded_sc.id}\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SingleCellTrajectory: saving & loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "SORT tracking...: 100%|██████████| 3/3 [00:00<00:00, 42.01it/s]\n"
     ]
    }
   ],
   "source": [
    "from typing import List\n",
    "from livecellx.track.sort_tracker_utils import (\n",
    "    gen_SORT_detections_input_from_contours,\n",
    "    update_traj_collection_by_SORT_tracker_detection,\n",
    "    track_SORT_bbox_from_contours,\n",
    "    track_SORT_bbox_from_scs\n",
    ")\n",
    "\n",
    "\n",
    "sct_collection = track_SORT_bbox_from_scs(single_cells, dic_dataset, mask_dataset=mask_dataset, max_age=1, min_hits=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "sct_collection.write_json(io_out_dir/\"sct_collection.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-----> json loaded from test_io_output/sct_collection.json\n",
      "|-----> Creating SingleCellTrajectoryCollection from json_dict...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<livecellx.core.single_cell.SingleCellTrajectoryCollection at 0x1c432c6a0>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sct_collection.load_from_json_file(io_out_dir/\"sct_collection.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<livecellx.core.datasets.LiveCellImageDataset at 0x1c4343ee0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sct_collection.get_all_trajectories()[0].img_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-----> json loaded from test_io_output/sct_collection.json\n",
      "|-----> Creating SingleCellTrajectoryCollection from json_dict...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<livecellx.core.single_cell.SingleCellTrajectoryCollection at 0x1c436ba30>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sct_collection.write_json(io_out_dir/\"sct_collection.json\", dataset_json_dir=io_out_dir/\"another_dataset_folder\")\n",
    "sct_collection.load_from_json_file(io_out_dir/\"sct_collection.json\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "livecell",
   "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.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}