worked on superanimal notebook

DeepLabCut · n-poulsen · Jul 19, 2024 · Apr 6, 2023 · Apr 20, 2023 · Apr 25, 2023
commit 830f4ac9d16f01af2b75890a7e2c433691e368a7
diff --git a/examples/COLAB/COLAB_Pytorch_SuperAnimal.ipynb b/examples/COLAB/COLAB_Pytorch_SuperAnimal.ipynb
@@ -22,7 +22,7 @@
     "\n",
     "# 🦄 SuperAnimal in DeepLabCut PyTorch! 🔥\n",
     "\n",
-    "This notebook demos how to use our SuperAnimal models within DLC3. Please read more in [Ye et al. Nature Communications 2024](https://www.nature.com/articles/s41467-024-48792-2) about the available SuperAnimal models, and follow along below!\n",
+    "This notebook demos how to use our SuperAnimal models within DeepLabCut 3.0! Please read more in [Ye et al. Nature Communications 2024](https://www.nature.com/articles/s41467-024-48792-2) about the available SuperAnimal models, and follow along below!\n",
     "\n",
     "### **Let's get going: install DeepLabCut into COLAB:**\n",
     "\n",
@@ -54,7 +54,7 @@
     "id": "5h0vq6E50Z4W"
    },
    "source": [
-    "### PLEASE, click \"restart runtime\" from the output above before proceeding!"
+    "**PLEASE, click \"restart runtime\" from the output above before proceeding!**"
    ]
   },
   {
@@ -73,6 +73,7 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "import zipfile\n",
     "from pathlib import Path\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
@@ -101,7 +102,7 @@
    "source": [
     "## Zero-shot Image Inference & Video Inference\n",
     "SuperAnimal models are foundation animal pose models. They can be used for zero-shot predictions without further training on the data.\n",
-    "In this section, we show how to use SuperAnimal models to predict pose from images (given an image folder) and output the predicted images (with pose) into another dest folder."
+    "In this section, we show how to use SuperAnimal models to predict pose from images (given an image folder) and output the predicted images (with pose) into another destination folder."
    ]
   },
   {
@@ -110,9 +111,9 @@
     "id": "FvFzntDMxPoL"
    },
    "source": [
-    "## Zero-shot image inference\n",
+    "### Zero-shot image inference\n",
     "\n",
-    "- If you have a single Image you want to test, upload it here!"
+    "If you have a single Image you want to test, upload it here!"
    ]
   },
   {
@@ -159,7 +160,7 @@
     "id": "Jashzdjb0Z4Y"
    },
    "source": [
-    "### Select a SuperAnimal name and corresponding model architecture\n",
+    "#### Select a SuperAnimal name and corresponding model architecture\n",
     "\n",
     "Check Our Docs on [SuperAnimals](https://github.com/DeepLabCut/DeepLabCut/blob/main/docs/ModelZoo.md) to learn more!"
    ]
@@ -197,7 +198,7 @@
    "outputs": [],
    "source": [
     "# Note you need to enter max_individuals correctly to get the correct number of predictions in the image.\n",
-    "superanimal_analyze_images(\n",
+    "_ = superanimal_analyze_images(\n",
     "    superanimal_name,\n",
     "    model_name,\n",
     "    image_path,\n",
@@ -213,7 +214,8 @@
    },
    "source": [
     "### Zero-shot Video Inference\n",
-    "- Without Video adaptation (faster, but not self-supervised fine-tuned on your data!)"
+    "\n",
+    "This can be done with or without video adaptation (faster, but not self-supervised fine-tuned on your data!)."
    ]
   },
   {
@@ -290,7 +292,9 @@
     "id": "Zv3v0QgSJNOg"
    },
    "source": [
-    "### Zero-shot Video Inference without video adaptation"
+    "#### Zero-shot Video Inference without video adaptation\n",
+    "\n",
+    "The labeled video (and pose predictions for the video) are saved in `\"/content/\"`, with the labeled video name being `{your_video_name}_superanimal_{superanimal_name}_hrnetw32_labeled.mp4`."
    ]
   },
   {
@@ -308,12 +312,12 @@
    },
    "outputs": [],
    "source": [
-    "video_inference_superanimal(\n",
+    "_ = video_inference_superanimal(\n",
     "    videos=video_path,\n",
     "    superanimal_name=f\"{superanimal_name}_{model_name}\",\n",
     "    video_adapt=False,\n",
     "    max_individuals=max_individuals,\n",
-    "    dest_folder=\"/content/\"\n",
+    "    dest_folder=\"/content/\",\n",
     ")"
    ]
   },
@@ -323,7 +327,9 @@
     "id": "Z8Z5GSti0Z4Z"
    },
    "source": [
-    "### Zero-shot Video Inference with video adaptation (unsupervised)"
+    "#### Zero-shot Video Inference with video adaptation (unsupervised)\n",
+    "\n",
+    "The labeled video (and pose predictions for the video) are saved in `\"/content/\"`, with the labeled video name being `{your_video_name}_superanimal_{superanimal_name}_hrnetw32_labeled_after_adapt.mp4`."
    ]
   },
   {
@@ -355,7 +361,7 @@
    },
    "outputs": [],
    "source": [
-    "video_inference_superanimal(\n",
+    "_ = video_inference_superanimal(\n",
     "    videos=[video_path],\n",
     "    superanimal_name=f\"{superanimal_name}_{model_name}\",\n",
     "    video_adapt=True,\n",
@@ -375,14 +381,16 @@
    },
    "source": [
     "## Training with SuperAnimal\n",
-    "In this section, we compare different ways to train the model, using and without using SuperAnimal.\n",
-    "You can compare the evaluation results and get a sense of each baseline.\n",
-    "We have following baselines:\n",
+    "\n",
+    "In this section, we compare different ways to train models in DeepLabCut 3.0, with or without using SuperAnimal-pretrained models.\n",
+    "You can compare the evaluation results and get a sense of each baseline. We have following baselines:\n",
     "\n",
     "- ImageNet transfer learning (training without superanimal)\n",
     "- SuperAnimal transfer learning (baseline 1)\n",
     "- SuperAnimal naive fine-tuning (baseline 2)\n",
-    "- SuperAnimal memory-replay fine-tuning (baseline3)"
+    "- SuperAnimal memory-replay fine-tuning (baseline3)\n",
+    "\n",
+    "This should be done on one of your DeepLabCut projects! If you don't have a DeepLabCut project that you can use SuperAnimal models with, you can always using the example openfield dataset [available in the DeepLabCut repository](https://github.com/DeepLabCut/DeepLabCut/tree/main/examples/openfield-Pranav-2018-10-30)! "
    ]
   },
   {
@@ -391,7 +399,9 @@
     "id": "L2wxevEn0Z4a"
    },
    "source": [
-    "#### Uploading your DLC project into Drive. Note you have to zip your DLC project and select the zipped file"
+    "#### Uploading your DeepLabCut project into Google Colab\n",
+    "\n",
+    "Note you have to zip your DLC project and select the zipped file."
    ]
   },
   {
@@ -416,6 +426,7 @@
     "    with zipfile.ZipFile(zip_file_path, \"r\") as zip_ref:\n",
     "        zip_ref.extractall(\"/content/dlc_project_folder\")\n",
     "\n",
+    "\n",
     "print(\"Contents of the extracted folder:\")\n",
     "extracted_files = os.listdir(\"/content/dlc_project_folder\")\n",
     "for file in extracted_files:\n",