Merge pull request #58 from dattalab/loader_refactor

Loader refactor
dattalab · Jun 5, 2023 · 4bb72ac · 4bb72ac
2 parents 09aa72c + 36fe073
commit 4bb72ac
Show file tree

Hide file tree

Showing 4 changed files with 242 additions and 102 deletions.
diff --git a/docs/keypoint_moseq_colab.ipynb b/docs/keypoint_moseq_colab.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "preliminary-agency",
    "metadata": {},
@@ -16,6 +17,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f6660e59",
    "metadata": {},
@@ -38,6 +40,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f8a52043",
    "metadata": {},
@@ -60,6 +63,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "890e3df1",
    "metadata": {},
@@ -87,6 +91,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "11ef323f",
    "metadata": {},
@@ -106,6 +111,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f65bd597",
    "metadata": {},
@@ -152,6 +158,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "gothic-viking",
    "metadata": {},
@@ -186,14 +193,18 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "phantom-dating",
    "metadata": {},
    "source": [
     "## Load data\n",
     "\n",
-    "Data can be loaded from DeepLabCut, SLEAP or from any another source as long as it has the following format, where K is the number of keypoints and D is 2 or 3\n",
-    "- `coordinates`: dict from session names to keypoint coordinate arrays of shape (T,K,D). Each key should start with its video name (e.g. `coordinates[\"experiment1_etc\"]` would correspond to `experiment1.avi`. In general this will already be true if importing from SLEAP or DeepLabCut).\n",
+    "## Load data\n",
+    "\n",
+    "We provide loading functions for DeepLabCut, SLEAP, anipose, and SLEAP-anipose. In every case, the keypoint data are loaded to the following pair of dictionaries:\n",
+    "\n",
+    "- `coordinates`: dict from session names to keypoint coordinate arrays of shape (T,K,D) where D is 2 or 3. Each key should start with its video name (e.g. `coordinates[\"experiment1_etc\"]` would correspond to `experiment1.avi`. In general this will already be true if importing from SLEAP or DeepLabCut).\n",
     "    \n",
     "- `confidences`: dict from session names to **nonnegative** keypoint confidence arrays of shape (T,K). Confidences are optional (they are used to set the error prior for each observation)."
    ]
@@ -205,19 +216,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# load data from DeepLabCut\n",
-    "dlc_results_directory = '/content/drive/MyDrive/dlc_project/videos/'\n",
-    "coordinates, confidences, bodyparts = kpms.load_deeplabcut_results(dlc_results_directory)\n",
-    "\n",
-    "# load data from SLEAP\n",
-    "# sleap_results_directory = '...'\n",
-    "# coordinates, confidences, bodyparts = kpms.load_sleap_results(sleap_results_directory)\n",
+    "# load data (e.g. from DeepLabCut)\n",
+    "keypoint_data_path = 'dlc_project/videos/' # can be a file, a directory, or a list of files\n",
+    "coordinates, confidences, bodyparts = kpms.load_keypoints(keypoint_data_path, 'deeplabcut')\n",
     "\n",
     "# format data for modeling\n",
     "data, labels = kpms.format_data(coordinates, confidences=confidences, **config())"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "processed-struggle",
    "metadata": {},
@@ -228,6 +236,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "organizational-theorem",
    "metadata": {},
@@ -269,6 +278,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "accomplished-pantyhose",
    "metadata": {},
@@ -294,6 +304,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "utility-penetration",
    "metadata": {},
@@ -316,6 +327,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "partial-remove",
    "metadata": {},
@@ -346,6 +358,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "thorough-identity",
    "metadata": {},
@@ -435,6 +448,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "37ae4b70",
    "metadata": {},
@@ -473,6 +487,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "breeding-fashion",
    "metadata": {},
@@ -481,6 +496,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "a2491a0d",
    "metadata": {},
@@ -500,6 +516,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "617a66ed",
    "metadata": {},

diff --git a/docs/source/FAQs.rst b/docs/source/FAQs.rst
@@ -72,7 +72,7 @@ Keypoint-MoSeq has only been validated on rodents (mice, rats, and anecdotal suc
 
 Loading data from methods other than SLEAP or DeepLabCut
 --------------------------------------------------------
-Keypoint-MoSeq can be used with any method that produces 2D or 3D keypoint detections. You can write a custom loading function or get in touch and request it as a new feature. 
+Keypoint-MoSeq can be used with any method that produces 2D or 3D keypoint detections. Currently we support SLEAP, DeepLabCut, anipose, and SLEAP-anipose (see :py:func:`keypoint_moseq.io.load_keypoints`). For methods not on this list, you can write a custom loading function or get in touch and request it as a new feature. 
 
 - If writing your own data loader, the output should be a ``coordinates`` dictionary that maps session names to arrays of shape ``(num_frames, num_keypoints, num_dimensions)``, where ``num_dimensions`` is 2 or 3. The keypoint axis should correspond to the `bodyparts` list in the config. You can also include a ``confidences`` dictionary that maps session names to arrays of shape ``(num_frames, num_keypoints)``. If your loader applies to a commonly used keypoint inference method, please let us know! We'd love to add it for others to use.
 

diff --git a/docs/source/tutorial.ipynb b/docs/source/tutorial.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "preliminary-agency",
    "metadata": {},
@@ -9,6 +10,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f8a52043",
    "metadata": {},
@@ -107,6 +109,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "gothic-viking",
    "metadata": {},
@@ -141,18 +144,20 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "phantom-dating",
    "metadata": {},
    "source": [
     "## Load data\n",
     "\n",
-    "Data can be loaded from DeepLabCut, SLEAP or from any another source as long as it has the following format, where K is the number of keypoints and D is 2 or 3\n",
-    "- `coordinates`: dict from session names to keypoint coordinate arrays of shape (T,K,D) \n",
+    "We provide loading functions for DeepLabCut, SLEAP, anipose, and SLEAP-anipose. In every case, the keypoint data are loaded to the following pair of dictionaries:\n",
+    "- `coordinates`: dict from session names to keypoint coordinate arrays of shape (T,K,2 [or 3]) \n",
     "- `confidences`: dict from session names to **nonnegative** keypoint confidence arrays of shape (T,K) "
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8ed16e05",
    "metadata": {},
@@ -172,19 +177,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# load data from DeepLabCut\n",
-    "dlc_results = 'dlc_project/videos/' # can be a file, a directory, or a list of files\n",
-    "coordinates, confidences, bodyparts = kpms.load_deeplabcut_results(dlc_results)\n",
-    "\n",
-    "# load data from SLEAP\n",
-    "# sleap_results = '...' # can be a file, a directory, or a list of files\n",
-    "# coordinates, confidences, bodyparts = kpms.load_sleap_results(sleap_results)\n",
+    "# load data (e.g. from DeepLabCut)\n",
+    "keypoint_data_path = 'dlc_project/videos/' # can be a file, a directory, or a list of files\n",
+    "coordinates, confidences, bodyparts = kpms.load_keypoints(keypoint_data_path, 'deeplabcut')\n",
     "\n",
     "# format data for modeling\n",
     "data, labels = kpms.format_data(coordinates, confidences=confidences, **config())"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "processed-struggle",
    "metadata": {},
@@ -421,6 +423,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "organizational-theorem",
    "metadata": {},
@@ -492,6 +495,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "accomplished-pantyhose",
    "metadata": {},
@@ -517,6 +521,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "utility-penetration",
    "metadata": {},
@@ -551,6 +556,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "partial-remove",
    "metadata": {},
@@ -699,6 +705,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "thorough-identity",
    "metadata": {},
@@ -839,6 +846,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "empty-houston",
    "metadata": {},
@@ -886,6 +894,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "a2491a0d",
    "metadata": {},
@@ -930,6 +939,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "617a66ed",
    "metadata": {},