UE5_parsers/Generate_DLC_Style_Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DeepLabCut style dataset parser\n",
    "\n",
    "### The generated output includes the following annotation data:\n",
    "* Classes\n",
    "* 2D keypoint data _(poses)_\n",
    "\n",
    "### Example application(s) (as demonstrated in Plum et al. 2023):\n",
    "* [DeepLabCut](https://www.mackenziemathislab.org/deeplabcut) _(single and multi-animal pose estimation)_\n",
    "\n",
    "### Output structure:\n",
    "* target_dir\n",
    "    * CollectedData_SCORER.csv _(pose annoation file)_\n",
    "    * CollectedData_SCORER.h5 _(pose annoation file)_\n",
    "    * all generated images\n",
    "    \n",
    "### Notes:\n",
    "\n",
    "* For ease of use, consider pointing the _**target_dir**_ to the _**\"labelled_data\"**_ directory in your created **DeepLabCut** project\n",
    "\n",
    "* **WARNING :** Restrict the **Colony size** to a **maxmimum of 20 individuals**, otherwise the generation of datasets will fail, as the number of necessary columns would exceed the 64kb Metadata limit of the exported HDF5 file!\n",
    "\n",
    "* Ensure, that your **GPU** has sufficient **memory** for the chosen resolution! E.g. at 1024 px X 1024 px, training on an RTX 2080 Ti, set the **batchsize** to 4 in the **pose_cfg.yaml** file of your model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import pathlib\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib as plt\n",
    "\n",
    "from os import listdir\n",
    "from os.path import isfile, join"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Required parameters\n",
    "\n",
    "Specify the location of your **generated dataset** and in which **output directory** you wish to save it.\n",
    "\n",
    "**Notes:**\n",
    "* do not include trailing forward slashes in your paths (see examples below)\n",
    "* Your **dataset** name should **NOT include underscores** as they are used to separate passes into their categories. Instead, use hyphens in your naming convention where applicable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define location of dataset and return all files\n",
    "\n",
    "# single-animal example\n",
    "dataset_location = \"../example_data/input-single\"\n",
    "target_dir = \"../example_data/DLC_SINGLE\"\n",
    "\n",
    "# multi-animal example\n",
    "#dataset_location = \"../example_data/input-multi\"\n",
    "#target_dir = \"../example_data/DLC_MULTI\"\n",
    "\n",
    "SCORER = \"Fabi\"\n",
    "\n",
    "# specify which labels to ignore. By default, all keypoints are written into the dataset\n",
    "# in this example we omit all keypoints relating to wings. Refer to the base_rig documentation for naming conventions\n",
    "omit_labels = ['w_1_l', 'w_1_l_end', 'w_2_l', 'w_2_l_end', 'w_1_r', 'w_1_r_end', 'w_2_r', 'w_2_r_end', 'root']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set True to show processing results for each image (disables parallel processing)\n",
    "DEBUG = False\n",
    "\n",
    "# we can optionally remove occluded points from the dataframe\n",
    "EXCLUDE_OCCLUDED_KEYPOINTS = True\n",
    "\n",
    "enforce_single_class = False # overwrites multiple classes and groups all instances as one"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following lines will load the generated dataset from your drive and prepare it for the multi-threaded parsing process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "all_files = [f for f in listdir(dataset_location) if isfile(join(dataset_location, f))]\n",
    "all_files.sort()\n",
    "\n",
    "# next, sort files into images, depth maps, segmentation maps, data, and colony info\n",
    "# we only need the location and name of the data files, as all passes follow the same naming convention\n",
    "dataset_data = []\n",
    "dataset_img = []\n",
    "dataset_ID = []\n",
    "dataset_depth = []\n",
    "dataset_norm = []\n",
    "dataset_colony = None\n",
    "\n",
    "for file in all_files:\n",
    "    loc = dataset_location + \"/\" + file\n",
    "    file_info = file.split(\"_\")\n",
    "    \n",
    "    if file_info[1] == \"BatchData\":\n",
    "        dataset_colony = loc\n",
    "        \n",
    "    elif len(file_info) == 2:\n",
    "        # images are available in various formats, but annotation data is always written as json files\n",
    "        if file_info[-1].split(\".\")[-1] == \"json\":\n",
    "            dataset_data.append(loc)\n",
    "        else:\n",
    "            dataset_img.append(loc)\n",
    "            \n",
    "    elif file_info[2].split(\".\")[0] == \"ID\":\n",
    "        dataset_ID.append(loc)\n",
    "    elif file_info[2].split(\".\")[0]  == \"depth\":\n",
    "        dataset_depth.append(loc)\n",
    "    elif file_info[2].split(\".\")[0]  == \"norm\":\n",
    "        dataset_norm.append(loc)\n",
    "        \n",
    "print(\"Found\",len(dataset_data),\"samples...\")\n",
    "\n",
    "# next sort the colony info into its IDs to determine the colony size and individual scales\n",
    "# Opening colony (BatchData) JSON file\n",
    "colony_file = open(dataset_colony)\n",
    " \n",
    "# returns JSON object as a dictionary\n",
    "colony = json.load(colony_file)\n",
    "colony_file.close()\n",
    "\n",
    "\n",
    "\"\"\" !!! requires IDs, model names, scales !!! \"\"\"\n",
    "\n",
    "\n",
    "if not enforce_single_class:\n",
    "    # get provided classes to create a dictionary of class IDs and class names\n",
    "    subject_class_names = np.unique(np.array(colony[\"Subject Variations\"]))\n",
    "    subject_classes = {}\n",
    "    for id,sbj in enumerate(subject_class_names):\n",
    "        subject_classes[str(sbj)] = id\n",
    "else:\n",
    "    subject_class_names = np.array([0])\n",
    "    subject_classes = {\"insect\" : 0}\n",
    "\n",
    "print(\"\\nA total of\",len(subject_class_names),\"unique classes have been found.\")\n",
    "print(\"The classes and respective class IDs are:\\n\",subject_classes,\"\\n\")\n",
    "\n",
    "\n",
    "print(\"Loaded colony file with seed\", colony['Seed']) #,\"and\",len(colony['ID']),\"individuals.\")\n",
    "    \n",
    "if len(colony['Subject Variations']) > 1:\n",
    "    multi_animal = True\n",
    "    print(\"Generating MULTI-animal dataset! Containing\",len(colony['Subject Variations']),\"individuals\")\n",
    "else:\n",
    "    multi_animal = False\n",
    "    print(\"Generating SINGLE-animal dataset!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As there may be animals for which we don't use all bones we can return a list of all labels and exclude the respective locations from the pose data. As all animals use the same convention, we can simply read in one example and remove the corresponding indices from all animals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading the first entry of first iteration file to retrieve skeleton info\n",
    "exp_file = open(dataset_data[0])\n",
    "exp_data = json.load(exp_file)\n",
    "exp_file.close()\n",
    "\n",
    "# for simplicity we'll assume that at this stage all subjects use the same armature and therefore report the same keypoints\n",
    "first_entry_key = list(exp_data[\"iterationData\"][\"subject Data\"][0].keys())[0]\n",
    "labels = list(exp_data[\"iterationData\"][\"subject Data\"][0][first_entry_key][\"keypoints\"].keys())\n",
    "\n",
    "# show all used labels:\n",
    "print(\"\\nAll labels:\",labels)\n",
    "\n",
    "print(\"\\nOmitting labels:\", omit_labels)\n",
    "\n",
    "# removing all occurences of omitted labels from the labels list to be used as keys below\n",
    "labels = [x for x in labels if x not in omit_labels]\n",
    "\n",
    "print(\"\\nFinal labels:\",labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us create a big array to store all our dataset info and then later convert it and save it all to the desired **.csv** and **.h5** files for DeepLabCut to read."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_points = np.zeros((len(dataset_data), (len(colony['Subject Variations'])*(len(labels)*2))))\n",
    "#\t- scorer   #(just one, the only scorer is the generator)\n",
    "#\t- - individuals\n",
    "#\t- - - bodyparts\n",
    "#\t- - - - coords\n",
    "\n",
    "print(\"Number of loaded samples:\",len(dataset_data))\n",
    "print(\"Colony size:\",len(colony['Subject Variations']))\n",
    "print(\"body parts:\",int(len(labels)),\"\\n\")\n",
    "print(\"Resulting in an array of shape:\",all_points.shape)\n",
    "\n",
    "output_file_names = [\"\" for i in range(len(dataset_data))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With all dataset related parameters configured, we have provided a multi-threaded parsing solution below to minimise the processing time it takes to bring the entire dataset into the required output format. Currently, we instanciate one processing thread per (virtual) CPU core but you can adjust this value if you wish by changing:\n",
    "\n",
    "```\n",
    "threadList_export = createThreadList(#NumDesiredThreads)\n",
    "```\n",
    "\n",
    "**Note:** To see the process of mask generation from ID passes in action, set the **DEBUG** mode to **\"True\"**. This will however slow down the processing speed considerably and only run in single-threaded mode!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# create unique colours for each ID\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "# alright. Let's take it from the top and fucking multi-thread this.\n",
    "import threading\n",
    "import queue\n",
    "import sys\n",
    "import os\n",
    "\n",
    "def fix_bounding_boxes(coords,max_val = [1024,1024]):\n",
    "    # fix bounding box coordinates so they do not reach beyond the image\n",
    "    fixed_coords = []\n",
    "    for c, coord in enumerate(coords):\n",
    "        if c == 0 or c == 2:\n",
    "            max_val_temp = max_val[0]\n",
    "        else:\n",
    "            max_val_temp = max_val[1]\n",
    "            \n",
    "        if coord >= max_val_temp:\n",
    "            coord = max_val_temp\n",
    "        elif coord <= 0:\n",
    "            coord = 0\n",
    "        \n",
    "        fixed_coords.append(int(coord))\n",
    "        \n",
    "    return fixed_coords\n",
    "\n",
    "def getThreads():\n",
    "    \"\"\" Returns the number of available threads on a posix/win based system \"\"\"\n",
    "    if sys.platform == 'win32':\n",
    "        return int(os.environ['NUMBER_OF_PROCESSORS'])\n",
    "    else:\n",
    "        return int(os.popen('grep -c cores /proc/cpuinfo').read())\n",
    "\n",
    "class exportThread(threading.Thread):\n",
    "    def __init__(self, threadID, name, q):\n",
    "        threading.Thread.__init__(self)\n",
    "        self.threadID = threadID\n",
    "        self.name = name\n",
    "        self.q = q\n",
    "\n",
    "    def run(self):\n",
    "        print(\"Starting \" + self.name)\n",
    "        process_detections(self.name, self.q)\n",
    "        print(\"Exiting \" + self.name)\n",
    "        \n",
    "def createThreadList(num_threads):\n",
    "    threadNames = []\n",
    "    for t in range(num_threads):\n",
    "        threadNames.append(\"Thread_\" + str(t))\n",
    "\n",
    "    return threadNames\n",
    "\n",
    "def process_detections(threadName, q):\n",
    "    while not exitFlag:\n",
    "        queueLock.acquire()\n",
    "        if not workQueue.empty():\n",
    "            \n",
    "            data_input = q.get()\n",
    "            i, data_loc, img, ID = data_input\n",
    "            queueLock.release()\n",
    "            \n",
    "            display_img = cv2.imread(img)\n",
    "            display_img_orig = display_img.copy()\n",
    "            \n",
    "            # compute visibility for each individual\n",
    "            seg_img = cv2.imread(ID)\n",
    "            seg_img_display = seg_img.copy()\n",
    "            \n",
    "            data_file = open(data_loc)\n",
    "            # returns JSON object as a dictionary\n",
    "            data = json.load(data_file)\n",
    "            data_file.close()\n",
    "            \n",
    "            img_shape = display_img.shape\n",
    "            \n",
    "            # only add images that contain visibile individuals\n",
    "            is_empty = True\n",
    "            \n",
    "            img_name = target_dir + \"/\" + img.split('/')[-1][:-4] + \"_synth\" + \".png\"\n",
    "            # write the file path to the all_points array\n",
    "            output_file_names[i] = \"labeled-data/\" + str(os.path.basename(target_dir)) + \"/\" + str(os.path.basename(img))[:-4] + \"_synth\" + \".png\"\n",
    "\n",
    "            img_info = []\n",
    "                \n",
    "            # check if the size of the image and segmentation pass match\n",
    "            if img_shape != seg_img.shape:\n",
    "                print(\"Size mismatch of image and segmentation pass for sample\",data_input[1].split(\"/\")[-1],\"!\")\n",
    "                incorrectly_formatted_images.append(i)\n",
    "            else:\n",
    "                for individual in data[\"iterationData\"][\"subject Data\"]:\n",
    "                    ind_key = list(individual.keys())[0]\n",
    "                    ind_ID = int(ind_key)\n",
    "                    # WARNING ID numbering begins at 1\n",
    "\n",
    "                    fontColor = (int(ID_colours[ind_ID,0]),\n",
    "                                 int(ID_colours[ind_ID,1]),\n",
    "                                 int(ID_colours[ind_ID,2]))\n",
    "                    \n",
    "                    bbox_orig = [individual[ind_key][\"2DBounds\"][\"xmin\"],\n",
    "                                 individual[ind_key][\"2DBounds\"][\"ymin\"],\n",
    "                                 individual[ind_key][\"2DBounds\"][\"xmax\"],\n",
    "                                 individual[ind_key][\"2DBounds\"][\"ymax\"]]\n",
    "                    \n",
    "                    bbox = fix_bounding_boxes(bbox_orig, max_val=display_img.shape)\n",
    "                    \n",
    "                    # only process an individual if its bounding box width and height are not zero\n",
    "                    if bbox[2] - bbox[0] == 0 or bbox[3] - bbox[1] == 0:\n",
    "                        continue\n",
    "\n",
    "                    try:\n",
    "                        ID_mask = cv2.inRange(seg_img[bbox[1]:bbox[3],bbox[0]:bbox[2]], np.array([0, 0, ind_ID - 2]), np.array([0, 0, ind_ID + 2]))\n",
    "                        indivual_occupancy = cv2.countNonZero(ID_mask)\n",
    "                    except:\n",
    "                        if len(threadList) == 1: \n",
    "                            print(\"Individual fully occluded:\",ind_ID,\"in\",dataset_seg[i])\n",
    "                        indivual_occupancy = 1\n",
    "\n",
    "                    #indivual_occupancy = np.count_nonzero((seg_img == [0, 0, int((individual[0]/len(colony['ID']))*255)]).all(axis = 2)) + np.count_nonzero((seg_img == [0, 0, int((individual[0]/len(colony['ID']))*255 - 1)]).all(axis = 2)) + np.count_nonzero((seg_img == [0, 0, int((individual[0]/len(colony['ID']))*255 + 1)]).all(axis = 2))\n",
    "                    bbox_area = abs((bbox[2] - bbox[0]) * (bbox[3] - bbox[1])) + 1\n",
    "                    bbox_occupancy = indivual_occupancy / bbox_area\n",
    "                    #print(\"Individual\", individual[0], \"with bounding box occupancy \",bbox_occupancy)\n",
    "\n",
    "                    #cv2.putText(display_img, \"ID: \" + str(int(individual[0])), (bbox[0] + 10,bbox[3] - 10), font, fontScale, fontColor, lineType)\n",
    "                    if bbox_occupancy > visibility_threshold:\n",
    "                        # let's binarise the image and dilate it to make sure all points that visible are found\n",
    "                        if not multi_animal:\n",
    "                            seg_bin = cv2.inRange(seg_img, np.array([0, 0, 1]), np.array([0,0, 3]))\n",
    "                        else:\n",
    "                            seg_bin = cv2.inRange(seg_img, np.array([0, 0, ind_ID - 1]), np.array([0, 0, ind_ID + 1]))\n",
    "                        kernel = np.ones((5,5), np.uint8)\n",
    "                        seg_bin_dilated = cv2.dilate(seg_bin,kernel,iterations = 2)\n",
    "                        if DEBUG:\n",
    "                            cv2.imshow(\"dilated mask\",seg_bin_dilated)\n",
    "                            cv2.waitKey(0)\n",
    "\n",
    "                        for point in range(len(labels)):\n",
    "                            # get rid of all invalid points first. Those should simply stay NaN in the array\n",
    "                            if individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] > img_shape[0] or individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] < 0 or individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] > img_shape[1] or individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] < 0:\n",
    "                                continue\n",
    "                            else:\n",
    "                                # now throw the coordinates to the correct location\n",
    "                                out_row = i\n",
    "                                out_column = ((ind_ID - 1) * (len(labels) ) + point) * 2\n",
    "                                # exclude negative keypoints\n",
    "                                if individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] < 0.1 or individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] < 0.1:\n",
    "                                    individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] = 0 # X\n",
    "                                    individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] = 0 # Y\n",
    "                                # exlucde occluded keypoints by checking their visibility in the segmentation map   \n",
    "                                if EXCLUDE_OCCLUDED_KEYPOINTS:\n",
    "                                    x_temp = int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"])\n",
    "                                    y_temp = int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"])           \n",
    "                                    if seg_bin_dilated[y_temp,x_temp] == 0:  \n",
    "                                        \n",
    "                                        if DEBUG:\n",
    "                                            display_img = cv2.circle(display_img, (x_temp,y_temp), radius=0, color=(0, 0, 255), thickness=2)\n",
    "                                            cv2.imshow(\"missing points\",display_img)\n",
    "                                            cv2.waitKey(0)\n",
    "                                        \n",
    "                                        individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] = 0 # X\n",
    "                                        individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] = 0 # Y\n",
    "                                all_points[out_row][out_column] = round(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"], 1) # X\n",
    "                                all_points[out_row][out_column + 1] = round(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] ,1) # Y\n",
    "\n",
    "                cv2.imwrite(img_name, display_img)\n",
    "            \n",
    "        else:\n",
    "            queueLock.release()\n",
    "            \n",
    "# setup as many threads as there are (virtual) CPU cores\n",
    "exitFlag = 0\n",
    "# only use a fourth of the number of CPUs for export as hugin and enfuse utilise multi core processing in part\n",
    "if DEBUG:\n",
    "    threadList = createThreadList(1)\n",
    "else:\n",
    "    threadList = createThreadList(getThreads())\n",
    "print(\"Using\", len(threadList), \"threads for export...\")\n",
    "queueLock = threading.Lock()\n",
    "\n",
    "# define paths to all images and set the maximum number of items in the queue equivalent to the number of images\n",
    "workQueue = queue.Queue(len(dataset_img))\n",
    "threads = []\n",
    "threadID = 1\n",
    "\n",
    "# keep track of all incorrectly formatted images to remove them after iterating over all entries\n",
    "incorrectly_formatted_images = []\n",
    "\n",
    "np.random.seed(seed=1)\n",
    "ID_colours = np.random.randint(255, size=(255, 3))\n",
    "\n",
    "font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "fontScale = 0.5\n",
    "lineType = 2\n",
    "\n",
    "# we can additionally plot the points in the data files to check joint locations\n",
    "plot_joints = True\n",
    "\n",
    "# remember to define an export folder when saving out your dataset\n",
    "generate_dataset = True\n",
    "\n",
    "# determine the proportion of a bounding box that needs to be filled before considering the visibility as too low\n",
    "# WARNING: At the moment the ID shown in segmentation maps does not always correspond to the ID in the data file (off by 1)\n",
    "visibility_threshold = 0.05\n",
    "\n",
    "timer = time.time()\n",
    "\n",
    "# Create new threads\n",
    "for tName in threadList:\n",
    "    thread = exportThread(threadID, tName, workQueue)\n",
    "    thread.start()\n",
    "    threads.append(thread)\n",
    "    threadID += 1\n",
    "\n",
    "# Fill the queue with samples\n",
    "queueLock.acquire()\n",
    "for i, (data, img, ID) in enumerate(zip(dataset_data , dataset_img, dataset_ID)):\n",
    "    workQueue.put([i, data, img, ID])\n",
    "queueLock.release()\n",
    "\n",
    "# Wait for queue to empty\n",
    "while not workQueue.empty():\n",
    "    pass\n",
    "\n",
    "# Notify threads it's time to exit\n",
    "exitFlag = 1\n",
    "\n",
    "# Wait for all threads to complete\n",
    "for t in threads:\n",
    "    t.join()\n",
    "print(\"Exiting Main export Thread\")\n",
    "\n",
    "# close all windows if they were opened\n",
    "cv2.destroyAllWindows()\n",
    "\n",
    "# now, remove all incorrectly formatted images from the points and file list\n",
    "all_points = np.delete(all_points, incorrectly_formatted_images ,axis=0)\n",
    "for r, rem_img in enumerate(incorrectly_formatted_images):\n",
    "    del output_file_names[rem_img - r]\n",
    "\n",
    "print(\"Total time elapsed:\",time.time()-timer,\"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, dump it all into one **DLC-conform pandas (.h5)** file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next create the required hierarchy\n",
    "scorer = [SCORER for i in range(len(all_points[0]))]\n",
    "individuals = []\n",
    "\n",
    "for ind in range(len(colony['Subject Variations'])):\n",
    "    ### UPDATE ONCE ALL COLONY INFO IS INCLUDED ###\n",
    "    individual = [\"id_\" + str(ind) +\"_num_\" + str(ind) for i in range(int((len(labels))*2))]\n",
    "    individuals.extend(individual)\n",
    "\n",
    "\n",
    "bodyparts_filtered = [i for j in labels for i in [j]*2]\n",
    "bodyparts = []\n",
    "\n",
    "for i in range(len(colony['Subject Variations'])):\n",
    "    bodyparts.extend(bodyparts_filtered)\n",
    "\n",
    "coords = []\n",
    "for i in range(int(len(all_points[0])/2)):\n",
    "    coords.extend([\"x\",\"y\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that all elemts for the **Multi-Index** hierachy are defined, we can combine them into the **final dataframe**  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if multi_animal:\n",
    "    categories = [scorer, individuals, bodyparts, coords]\n",
    "    categories_tuples = list(zip(*categories))\n",
    "    columns = pd.MultiIndex.from_tuples(categories_tuples, names=[\"scorer\",\n",
    "                                                               \"individuals\",\n",
    "                                                               \"bodyparts\",\n",
    "                                                               \"coords\"])\n",
    "else:\n",
    "    categories = [scorer, bodyparts, coords]\n",
    "    categories_tuples = list(zip(*categories))\n",
    "    columns = pd.MultiIndex.from_tuples(categories_tuples, names=[\"scorer\",\n",
    "                                                               \"bodyparts\",\n",
    "                                                               \"coords\"])\n",
    "    \n",
    "final_dataframe = pd.DataFrame(all_points, index = output_file_names, columns=columns)\n",
    "# convert all zeros to NaN\n",
    "final_dataframe = final_dataframe.replace(0, np.nan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Display the output dataframe..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "final_dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and write it all out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_dataframe.to_csv(os.path.join(target_dir, \"CollectedData_\" + SCORER + \".csv\"))\n",
    "\n",
    "# IF the function below fails, this is likely due to exceeding the number of columns supported by HDF5 files!\n",
    "# Restrict the number of simulated animals to < 20 if the goal is to train a DLC network\n",
    "\n",
    "final_dataframe.to_hdf(\n",
    "    os.path.join(target_dir, \"CollectedData_\" + SCORER + \".h5\"),\n",
    "    \"df_with_missing\",\n",
    "    format=\"table\",\n",
    "    mode=\"w\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}