updated colab to match modeling.ipynb

Author: JRJacoby

docs/keypoint_moseq_colab.ipynb

@@ -33,9 +33,10 @@
     "! pip install -U keypoint-moseq\n",
     "\n",
     "import os\n",
-    "from google.colab import drive\n",
+    "from google.colab import drive, output\n",
     "\n",
-    "drive.mount(\"/content/drive\")"
+    "drive.mount(\"/content/drive\")\n",
+    "output.enable_custom_widget_manager()"
    ]
   },
   {
@@ -277,7 +278,7 @@
     "Removing large outliers can improve the robustness of model fitting. The following cell classifies keypoints as outliers based on their distance to the animal's medoid. The outlier keypoints are then interpolated and their confidences are set to 0.\n",
     "- Use `outlier_scale_factor` to adjust the stringency of outlier detection (higher values -> more stringent)\n",
     "- Plots showing distance to medoid before and after outlier interpolation are saved to `{project_dir}/QA/plots/`\n",
-    "- Plotting can take a few minutes. If you have a lot of data, you can comment out the plot_medoid_distance_outliers line."
+    "- Plotting can take a few minutes, so by default plots will not be regenerated when re-running this cell. To experiment with the effects of setting different values for outlier_scale_factor, set `overwrite=True` in outlier_removal."
    ]
   },
   {
@@ -287,25 +288,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# outlier_scale_factor modifies the how stringent the outlier detection algorithm is. Run this cell once and examine the output\n",
-    "# plots in {project_dir}/QA/plots/keypoint_distance_outliers. If not enough outliers are being interpolated, set this number\n",
-    "# lower. If too many correct points are being interpolated, set this number higher.\n",
     "kpms.update_config(project_dir, outlier_scale_factor=6.0)\n",
     "\n",
-    "for recording_name in coordinates:\n",
-    "    raw_coords = coordinates[recording_name].copy()\n",
-    "    outliers = kpms.find_medoid_distance_outliers(raw_coords, **config())\n",
-    "    coordinates[recording_name] = kpms.interpolate_keypoints(raw_coords, outliers[\"mask\"])\n",
-    "    confidences[recording_name] = np.where(outliers[\"mask\"], 0, confidences[recording_name])\n",
-    "    kpms.plot_medoid_distance_outliers(\n",
-    "        project_dir,\n",
-    "        recording_name,\n",
-    "        raw_coords,\n",
-    "        coordinates[recording_name],\n",
-    "        outliers[\"mask\"],\n",
-    "        outliers[\"thresholds\"],\n",
-    "        **config()\n",
-    "    )"
+    "coordinates, confidences = kpms.outlier_removal(\n",
+    "    coordinates,\n",
+    "    confidences,\n",
+    "    project_dir,\n",
+    "    overwrite=False,\n",
+    "    **config()\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e09604c",
+   "metadata": {},
+   "source": [
+    "## Format data for modeling"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c70fc436",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data, metadata = kpms.format_data(coordinates, confidences, **config())"
    ]
   },
   {
@@ -634,6 +643,13 @@
     "# # load new data (e.g. from deeplabcut)\n",
     "# new_data = 'path/to/new/data/' # can be a file, a directory, or a list of files\n",
     "# coordinates, confidences, bodyparts = kpms.load_keypoints(new_data, 'deeplabcut')\n",
+    "# coordinates, confidences = kpms.outlier_removal(\n",
+    "#     coordinates,\n",
+    "#     confidences,\n",
+    "#     project_dir,\n",
+    "#     overwrite=False,\n",
+    "#     **config()\n",
+    "# )\n",
     "# data, metadata = kpms.format_data(coordinates, confidences, **config())\n",
     "\n",
     "# # apply saved model to new data\n",