{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c1179595",
   "metadata": {},
   "source": [
    "# Using 2D skeletonisation to improve the length measurements of the leaves in a 3D maize plant segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07b7872",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from openalea.phenomenal.tracking.leaf_extension import (\n",
    "    skeleton_branches,\n",
    "    compute_extension,\n",
    "    leaf_extension,\n",
    ")\n",
    "\n",
    "import openalea.phenomenal.object.voxelSegmentation as phm_seg\n",
    "from openalea.phenomenal.calibration import Calibration\n",
    "\n",
    "datadir = \"./data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0d6a04",
   "metadata": {},
   "source": [
    "## 1. Load data : 12 binary images, 3D segmentation, 3D->2D projection functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79750ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "angles = [a * 30 for a in range(12)]\n",
    "\n",
    "binaries = {\n",
    "    angle: cv2.imread(datadir + \"/leaf_extension/{}.png\".format(angle), 0)\n",
    "    for angle in angles\n",
    "}\n",
    "\n",
    "# 3D segmentation\n",
    "seg = phm_seg.VoxelSegmentation.read_from_json_gz(\n",
    "    datadir + \"/leaf_extension/segmentation.gz\"\n",
    ")\n",
    "\n",
    "# 3D to 2D projection functions\n",
    "calibration = Calibration.load(datadir + \"/leaf_extension/calibration.json\")\n",
    "projections = {\n",
    "    angle: calibration.get_projection(\n",
    "        id_camera=\"side\", rotation=angle, world_frame=\"pot\"\n",
    "    )\n",
    "    for angle in angles\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a24792e2",
   "metadata": {},
   "source": [
    "## 2. Leaf extension on a single camera angle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c161c6c5",
   "metadata": {},
   "source": [
    "### 2D skeletonisation of the binary image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d721d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_angle = 90\n",
    "\n",
    "binary = binaries[example_angle]\n",
    "polylines_2d = skeleton_branches(binary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7640f737",
   "metadata": {},
   "source": [
    "### Display it\n",
    "\n",
    "+ From the different branches found in the 2D skeleton, only the ones which have an endpoint in the skekeleton graph are saved.\n",
    "+ These branches are displayed with a blue line below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff48531a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 10), dpi=100)\n",
    "plt.imshow(255 - binary, cmap=\"Greys\")\n",
    "for pl in polylines_2d:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"b-\")\n",
    "    plt.plot(pl[-1, 0], pl[-1, 1], \"bo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4aebd4",
   "metadata": {},
   "source": [
    "### Project leaf polylines from the 3D segmentation in 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f544130a",
   "metadata": {},
   "outputs": [],
   "source": [
    "polylines_3d = [\n",
    "    seg.get_leaf_order(k).real_longest_polyline()\n",
    "    for k in range(1, 1 + seg.get_number_of_leaf())\n",
    "]\n",
    "\n",
    "polylines_3d_to_2d = [projections[example_angle](pl) for pl in polylines_3d]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "200bc69a",
   "metadata": {},
   "source": [
    "### Show the comparison\n",
    "\n",
    "+ 3D skeletonisation and segmentation from Phenomenal (red lines below) finds more leaves, and manages leaf overlaps better. However, It sometimes loses the thin ends of the leaves.\n",
    "+ 2D skeletonisation (blue lines below) finds less leaves, but it is better at reaching the ends of the leaves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e30b34b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 10), dpi=100)\n",
    "plt.imshow(255 - binary, cmap=\"Greys\")\n",
    "for pl in polylines_2d:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"b-\")\n",
    "    plt.plot(pl[-1, 0], pl[-1, 1], \"bo\")\n",
    "for pl in polylines_3d_to_2d:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"r-\")\n",
    "    plt.plot(pl[-1, 0], pl[-1, 1], \"ro\")\n",
    "plt.title(\n",
    "    \"Blue : polylines from the 2D skeleton\"\n",
    "    + \"\\nRed : polylines from the 3D skeleton (phenomenal) projected in 2D\"\n",
    "    + \"\\n(Circles = end extremities)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0bc5e56",
   "metadata": {},
   "source": [
    "## Compute leaf extension\n",
    "This function tries to match both types of polylines. For each match, an extension path is determined, to correct the length measurement for the 3D leaf polyline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbd818c",
   "metadata": {},
   "outputs": [],
   "source": [
    "extension_factors, extension_polylines = compute_extension(\n",
    "    polylines_3d_to_2d, polylines_2d\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7855b7b",
   "metadata": {},
   "source": [
    "## Display it\n",
    "Extension paths are shown in green below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e9df02",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 10), dpi=100)\n",
    "plt.imshow(255 - binary, cmap=\"Greys\")\n",
    "for pl in polylines_2d:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"b-\")\n",
    "    plt.plot(pl[-1, 0], pl[-1, 1], \"bo\")\n",
    "for pl in polylines_3d_to_2d:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"r-\")\n",
    "    plt.plot(pl[-1, 0], pl[-1, 1], \"ro\")\n",
    "for pl in extension_polylines:\n",
    "    plt.plot(pl[:, 0], pl[:, 1], \"-\", color=\"lime\", linewidth=3)\n",
    "plt.title(\n",
    "    \"Blue : polylines from the 2D skeleton\"\n",
    "    + \"\\nRed : polylines from the 3D skeleton (phenomenal) projected in 2D\"\n",
    "    + \"\\nGreen : extension polylines\"\n",
    "    + \"\\n(Circles = end extremities)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8d28bc",
   "metadata": {},
   "source": [
    "## 3. Combine leaf extensions from multiple angles to improve the leaf length measurements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "908e417b",
   "metadata": {},
   "source": [
    "### Multi angle leaf extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b65e0f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_seg = leaf_extension(seg, binaries, projections)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0433125",
   "metadata": {},
   "source": [
    "### Compare the old and new length values for each leaf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a577f108",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = []\n",
    "for leaf in new_seg.get_leafs():\n",
    "    old_length = (\n",
    "        leaf.info[\"pm_length_with_speudo_stem\"]\n",
    "        if leaf.info[\"pm_label\"] == \"growing_leaf\"\n",
    "        else leaf.info[\"pm_length\"]\n",
    "    )\n",
    "    new_length = leaf.info[\"pm_length_extended\"]\n",
    "    df.append([old_length, new_length])\n",
    "df = pd.DataFrame(df, columns=[\"old_length\", \"new_length\"])\n",
    "print(df)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}