{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wormpose\n",
    "wormpose.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download utility function to display images in jupyter notebook\n",
    "!wget https://raw.githubusercontent.com/iteal/wormpose/master/examples/ipython_utils.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wormpose.dataset.loader import load_dataset\n",
    "from wormpose.dataset.loaders.resizer import add_resizing_arguments, ResizeOptions\n",
    "from wormpose.images.real_dataset import RealDataset\n",
    "\n",
    "class HeadPosDataset(object):\n",
    "    # this class only uses one video from the dataset, to extend for more\n",
    "    # it will  generate pairs of square images of worms + associated head position\n",
    "    # you can use resize options to resize the images (ex pass: image_size=100)\n",
    "    \n",
    "    def __init__(self, dataset_loader: str, dataset_path: str, video_name=None, **kwargs):        \n",
    "         \n",
    "        resize_options = ResizeOptions(**kwargs)\n",
    "        dataset = load_dataset(dataset_loader, dataset_path, resize_options=resize_options)\n",
    "        self.video_name = video_name if video_name is not None else dataset.video_names[0]\n",
    "        \n",
    "        self.real_dataset = RealDataset(\n",
    "            frame_preprocessing=dataset.frame_preprocessing,\n",
    "            output_image_shape=dataset.image_shape,\n",
    "        )        \n",
    "        self.frames_dataset = dataset.frames_dataset\n",
    "        # the head position is the first joint of the skeleton\n",
    "        self.head_positions = dataset.features_dataset[self.video_name].skeletons[:, 0,:]\n",
    "  \n",
    "    def generate(self):\n",
    "        with self.frames_dataset.open(self.video_name) as frames:\n",
    "            for i, frame in enumerate(frames):\n",
    "                processed_frame, crop_offset = self.real_dataset.process_frame(frame)\n",
    "                # convert the head position in crop coordinates\n",
    "                yield processed_frame, self.head_positions[i] - crop_offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipython_utils import ImagesViewer, display_as_slider\n",
    "import math\n",
    "\n",
    "dataset_loader = \"tierpsy\"\n",
    "# set the path to your tierpsy dataset\n",
    "# see https://iteal.github.io/wormpose/tierpsy.html for how to download and setup folders\n",
    "dataset_path = \"MY_PATH/TO/TIERPSY_DATASET\" \n",
    "\n",
    "head_pos_dataset = HeadPosDataset(dataset_loader, dataset_path).generate()\n",
    "processed_img_viewer = ImagesViewer()\n",
    "\n",
    "max_viz = 10\n",
    "print(f\"Displaying the first {max_viz} frames\")\n",
    "\n",
    "for _ in range(max_viz):\n",
    "    # example of iteration\n",
    "    processed_image, head_pos = next(head_pos_dataset)\n",
    "    \n",
    "    # here we draw on the image for visualization of the head\n",
    "    if not math.isnan(head_pos[0]):\n",
    "        cv2.circle(processed_image, (int(head_pos[0]), int(head_pos[1])), 3, 255, -1)\n",
    "        \n",
    "    processed_img_viewer.add_image(processed_image)\n",
    "\n",
    "\n",
    "display_as_slider(processed_img_viewer)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}