UE5_parsers/Generate_Custom_3D_Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom 3D style dataset parser\n",
    "\n",
    "### The generated output includes the following annotation data:\n",
    "*     bounding_box &nbsp;&nbsp;&nbsp; 4 x 1 float\n",
    "*\t  key_points_3D &nbsp;&nbsp;&nbsp;3 x k float (provide name sheet)\n",
    "*\t  key_points_2D &nbsp;&nbsp;&nbsp;2 x k float\n",
    "*\t  visibility &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 1 x k int (0 occluded or 1 visible)\n",
    "*\t  rot_mat\t&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 3 x 3 float\n",
    "*\t  trans_mat\t&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 3 x 1 float\n",
    "*\t  intrinsics_mat &nbsp;&nbsp;&nbsp; 3 x 3 float\n",
    "\n",
    "### Example application(s) (as demonstrated in Plum et al. 2023):\n",
    "* _experimental feature / not used in presented results_\n",
    "\n",
    "### Output structure:\n",
    "* target_dir\n",
    "    * Data_3D_Pose.hdf5 _(pose annoation file)_\n",
    "    * label_names.txt\n",
    "    * all generated images\n",
    "    \n",
    "### Notes:\n",
    "\n",
    "* This script is used to automatically generate custom datasets to support object detection as well as 3D and 2D pose estimation, including camera data.\n",
    "\n",
    "* **WARNING :** In this version, restrict the **Colony size** to a **maxmimum of 1 indivudal**!\n",
    "\n",
    "* The script **excludes empty samples** (*no animals present*) automatically and provide an additional **occlusion vector**, indicating whether a key point is visible (1) or occluded (0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import json\n",
    "import time\n",
    "import threading\n",
    "import queue\n",
    "import sys\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib as plt\n",
    "import pathlib\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",
    "dataset_location = \"../example_data/input-single\"\n",
    "target_dir = \"../example_data/3D\"\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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",
    "\"\"\" !!! requires IDs, model names, scales !!! \"\"\"\n",
    "\n",
    "print(\"Loaded colony file with seed\",colony['Seed'],\"and\",len(colony['Subject Variations']),\"individual(s).\")\n",
    "if len(colony['Subject Variations']) == 1:\n",
    "    print(\"Generating single-animal dataset!\")\n",
    "else:\n",
    "    print(\"WARNING! Multi-animal datasets are currently NOT supported!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the cleaned colony info, we can start loading the data associated with each frame.\n",
    "For simplicity, we will produce a list of lists, containing all individuals and their attributes for each frame.\n",
    "\n",
    "We will therefore access \"data\" as **[iteration] [individual] [attribute]**, where attributes will include [ID,bbox_x_0,bbox_y_0,...]\n",
    "\n",
    "The data files additionally contain **camera information**, such as **extrinsics** in the form of **transformation** and **rotation** in the global coordinate frame, and camera **intrinsics**. For ease of use, the data files also contain the camera **View Projection Matrix**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_data = []\n",
    "camera_data_types = []\n",
    "\n",
    "# loading the first iteration data file to retrieve keys of camera info\n",
    "exp_file = open(dataset_data[0])\n",
    "exp_data = json.load(exp_file)\n",
    "exp_file.close()\n",
    "\n",
    "# get the types of camera information stored\n",
    "camera_data_types = list(exp_data[\"iterationData\"][\"camera\"].keys())\n",
    "\n",
    "print(\"The following camera data has been included:\\n\",camera_data_types)"
   ]
  },
  {
   "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",
    "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 that loaded example annotation data and the batch / colony info we can start plotting bounding boxes, joint locations, and check if the camera attributes have been exported correctly.\n",
    "\n",
    "Let's quickly define a few functions to parse the produced data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform between sRGB and linear colour space (optional)\n",
    "\n",
    "def to_linear(srgb):\n",
    "    linear = np.float32(srgb) / 255.0\n",
    "    less = linear <= 0.04045\n",
    "    linear[less] = linear[less] / 12.92\n",
    "    linear[~less] = np.power((linear[~less] + 0.055) / 1.055, 2.4)\n",
    "    return linear * 255.0\n",
    "\n",
    "    \n",
    "def from_linear(linear):\n",
    "    srgb = linear.copy()\n",
    "    less = linear <= 0.0031308\n",
    "    srgb[less] = linear[less] * 12.92\n",
    "    srgb[~less] = 1.055 * np.power(linear[~less], 1.0 / 2.4) - 0.055\n",
    "    return srgb * 255.0\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",
    "# and compute the XYZ rotation matrix from roll, pitch, and yaw\n",
    "def get_rotation_matrix(roll,pitch,yaw,degrees=True):\n",
    "    # convert to radian\n",
    "    if degrees:\n",
    "        roll = np.radians(-roll)\n",
    "        pitch = np.radians(pitch)\n",
    "        yaw = np.radians(-yaw)\n",
    "    # roll rotation \n",
    "    Rx = np.array([[1, 0, 0],\n",
    "                   [0,np.cos(roll),-np.sin(roll)],\n",
    "                   [0,np.sin(roll),np.cos(roll)]])\n",
    "    # pitch rotation\n",
    "    Ry = np.array([[np.cos(pitch),0,np.sin(pitch)],\n",
    "                   [0, 1, 0],\n",
    "                   [-np.sin(pitch),0,np.cos(pitch)]])\n",
    "    # yaw rotation\n",
    "    Rz = np.array([[np.cos(yaw),-np.sin(yaw),0],\n",
    "                   [np.sin(yaw),np.cos(yaw),0],\n",
    "                   [0, 0, 1]])\n",
    "    #Rxyz = np.round(np.matmul(np.matmul(Rz,Ry),Rx),3)\n",
    "    Rxyz = Ry @ Rz @ Rx\n",
    "    return Rxyz\n",
    "\n",
    "def parse_projection_components(iteration_data_file):\n",
    "    \n",
    "    ####### DOUBLE CHECK THESE VALUES GO INTO THE RIGHT MATRIX ELEMENTS ############\n",
    "    \n",
    "    # converts Unreal view projection into rotation and translation components\n",
    "    input_matrix = iteration_data_file[\"iterationData\"][\"camera\"][\"View Matrix\"]\n",
    "    w = input_matrix[\"wPlane\"]\n",
    "    x = input_matrix[\"xPlane\"]\n",
    "    y = input_matrix[\"yPlane\"]\n",
    "    z = input_matrix[\"zPlane\"]\n",
    "    # now, assign the respective transposed values to the rotation...\n",
    "    cam_rot = np.array([[x[\"x\"],y[\"x\"],z[\"x\"]],\n",
    "                        [x[\"y\"],y[\"y\"],z[\"y\"]],\n",
    "                        [x[\"z\"],y[\"z\"],z[\"z\"]]])\n",
    "    # and the translation\n",
    "    \"\"\"\n",
    "    cam_trans = np.array([iteration_data_file[\"iterationData\"][\"camera\"][\"Location\"][\"x\"],\n",
    "                          iteration_data_file[\"iterationData\"][\"camera\"][\"Location\"][\"y\"],\n",
    "                          iteration_data_file[\"iterationData\"][\"camera\"][\"Location\"][\"z\"]])\n",
    "    \"\"\"\n",
    "    cam_trans = np.array([w[\"x\"],\n",
    "                          w[\"y\"],\n",
    "                          w[\"z\"]])\n",
    "    # There. Tried to do it differently, had a break down, now it works. \n",
    "    # Bon appetit\n",
    "    return cam_rot,cam_trans\n",
    "\n",
    "def parse_camera_intrinsics(batch_data_file, iteration_data_file):\n",
    "    # first get the image resolution from the batch data file and the current FOV from the iteration data file\n",
    "    res_px_X = batch_data_file[\"Image Resolution\"][\"x\"]\n",
    "    res_px_Y = batch_data_file[\"Image Resolution\"][\"y\"]\n",
    "    FOV = iteration_data_file[\"iterationData\"][\"camera\"][\"FOV\"]\n",
    "    \n",
    "    # then compute the image centre and focal length in x and y respectively\n",
    "    # https://docs.opencv.org/2.4/modules/calib3d/doc/camera_calibration_and_3d_reconstruction.html \n",
    "    \n",
    "    cx = res_px_X / 2\n",
    "    cy = res_px_Y / 2\n",
    "    \n",
    "    fx = cx / np.tan(np.radians(FOV)/2)\n",
    "    fy = cy / np.tan(np.radians(FOV)/2)\n",
    "    \n",
    "    return cx, cy, fx, fy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating 3D pose output files\n",
    "Now comes the difficult part: getting all this data into the required format.\n",
    "\n",
    "We're going to want an **.h5** formatted file, essentially one dataframe for the entire dataset with the following entries:\n",
    "\n",
    "*\t  file_name &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 1 string (relative)\n",
    "*\t  rot_mat\t&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 3 x 3 float\n",
    "*\t  trans_mat\t&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 3 x 1 float\n",
    "*\t  intrinsics_mat &nbsp;&nbsp;&nbsp; 3 x 3 float\n",
    "*     bounding_box &nbsp;&nbsp;&nbsp; 4 x 1 float\n",
    "*\t  key_points_3D &nbsp;&nbsp;&nbsp;3 x k float (provide name sheet)\n",
    "*\t  key_points_2D &nbsp;&nbsp;&nbsp;2 x k float\n",
    "*\t  visibility &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; 1 x k int (0 occluded or 1 visible)\n",
    "\n",
    "To provide visibility info, we will check whether the subject is visbile in the respective segmentation map at the given screen X & Y coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_df = pd.DataFrame(index=range(len(dataset_data)),columns=[\"file_name\",\n",
    "                                                      \"cam_rot\",\n",
    "                                                      \"cam_trans\",\n",
    "                                                      \"cam_intrinsics\",\n",
    "                                                      \"bounding_box\",\n",
    "                                                      \"key_points_3D\",\n",
    "                                                      \"key_points_2D\",\n",
    "                                                      \"visibility\"])\n",
    "\n",
    "print(\"Number samples:\",len(dataset_data))\n",
    "print(\"Colony size:\",len(colony['Subject Variations']))\n",
    "print(\"body parts:\",len(labels),\" (including image X & Y, as well as world X Y Z coordinates)\\n\")\n",
    "print(\"Resulting in a dataframe of shape:\",out_df.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": true
   },
   "outputs": [],
   "source": [
    "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_export:\n",
    "        queueLock.acquire()\n",
    "        if not workQueue_export.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] = 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",
    "                    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 DEBUG: \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",
    "                    # write all camera attributes to the output file\n",
    "                    # for details, refer to https://ksimek.github.io/2013/08/13/intrinsic/\n",
    "\n",
    "                    out_df.loc[i][\"cam_rot\"],out_df.loc[i][\"cam_trans\"] = parse_projection_components(data)\n",
    "\n",
    "                    cx, cy, fx, fy = parse_camera_intrinsics(colony,data)\n",
    "                    # https://docs.opencv.org/2.4/modules/calib3d/doc/camera_calibration_and_3d_reconstruction.html\n",
    "                    out_df.loc[i][\"cam_intrinsics\"] = np.array([[fx, 0,  cx],\n",
    "                                                                [0,  fy, cy],\n",
    "                                                                [0,  0,  1]])\n",
    "\n",
    "                    visbility_img = []\n",
    "                    XY_2D_points = []\n",
    "                    XYZ_3D_points = []\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",
    "                        seg_bin = cv2.inRange(seg_img, np.array([0, 0, 1]), np.array([0,0, 3]))\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",
    "\n",
    "                            if individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"] < 0.1 or individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"] < 0.1:\n",
    "                                # exclude all points that lie outside the image\n",
    "                                visibility_point = 0\n",
    "                                XY_2D_points.append([0,0])\n",
    "                                XYZ_3D_points.append([0,0,0])\n",
    "                            else:\n",
    "                                # check if the 2D point is occluded in the segmentation image\n",
    "                                # thanks opencv, of course this has to be indexed as Y,X... thanks, really.\n",
    "                                if seg_bin_dilated[int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"]),\n",
    "                                                   int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"])] == 255:                   \n",
    "                                    visibility_point = 1\n",
    "                                    XY_2D_points.append([individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"],\n",
    "                                                         individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"]])\n",
    "                                    XYZ_3D_points.append([individual[ind_key][\"keypoints\"][labels[point]][\"3DPos\"][\"x\"],\n",
    "                                                          individual[ind_key][\"keypoints\"][labels[point]][\"3DPos\"][\"y\"],\n",
    "                                                          individual[ind_key][\"keypoints\"][labels[point]][\"3DPos\"][\"z\"]])\n",
    "\n",
    "                                    # draw 2D points for visualisation\n",
    "                                    if DEBUG:\n",
    "                                        seg_img_display = cv2.circle(seg_img_display, (int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"x\"]),\n",
    "                                                                                       int(individual[ind_key][\"keypoints\"][labels[point]][\"2DPos\"][\"y\"])), \n",
    "                                                                                       radius=0, color=(255, 100, 100), \n",
    "                                                                                       thickness=5)\n",
    "                                else:\n",
    "                                    visibility_point = 0\n",
    "                                    XY_2D_points.append([0,0])\n",
    "                                    XYZ_3D_points.append([0,0,0])\n",
    "\n",
    "                            visbility_img.append(visibility_point)\n",
    "\n",
    "                    if len(threadList_export) == 1:       \n",
    "                        cv2.imshow(\"Segmentation and points\", seg_img_display)\n",
    "                        cv2.imshow(\"Segmentation binarised and dilated\", seg_bin_dilated)\n",
    "                        cv2.waitKey(0) \n",
    "\n",
    "                    out_df.loc[i][\"visibility\"] = visbility_img\n",
    "\n",
    "                # if no entries were found for the respective image, remove it from the output list and don't write out the image\n",
    "                if all([ v == 0 for v in out_df.loc[i][\"visibility\"]]):\n",
    "                    remove_empty_entries.append(i)\n",
    "                else:    \n",
    "                    out_df.loc[i][\"file_name\"] = output_file_names[i]\n",
    "                    out_df.loc[i][\"bounding_box\"] = bbox\n",
    "                    out_df.loc[i][\"key_points_2D\"] = XY_2D_points\n",
    "                    out_df.loc[i][\"key_points_3D\"] = XYZ_3D_points\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_export = 0\n",
    "# set the following to 1 (instead of getThreads()) to display segmentation maps and visible key points\n",
    "if DEBUG:\n",
    "    threadList_export = createThreadList(1)\n",
    "else:\n",
    "    threadList_export = createThreadList(getThreads())\n",
    "print(\"Using\", len(threadList_export), \"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_export = queue.Queue(len(dataset_img))\n",
    "threads = []\n",
    "threadID = 1\n",
    "\n",
    "np.random.seed(seed=1)\n",
    "\n",
    "font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "fontScale = 0.5\n",
    "lineType = 2\n",
    "\n",
    "# remove entries that contain no data\n",
    "remove_empty_entries = []\n",
    "\n",
    "# set true if generating dataset for animal without wings\n",
    "exclude_wings = True\n",
    "\n",
    "# we can additionally plot the points in the data files to check joint locations\n",
    "plot_joints = True\n",
    "\n",
    "# remember to refine 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.001\n",
    "\n",
    "timer = time.time()\n",
    "\n",
    "# Create new threads\n",
    "for tName in threadList_export:\n",
    "    thread = exportThread(threadID, tName, workQueue_export)\n",
    "    thread.start()\n",
    "    threads.append(thread)\n",
    "    threadID += 1\n",
    "\n",
    "# Fill the queue with stacks\n",
    "queueLock.acquire()\n",
    "for i, (data, img, ID) in enumerate(zip(dataset_data , dataset_img, dataset_ID)):\n",
    "    workQueue_export.put([i, data, img, ID])\n",
    "queueLock.release()\n",
    "\n",
    "# Wait for queue to empty\n",
    "while not workQueue_export.empty():\n",
    "    pass\n",
    "\n",
    "# Notify threads it's time to exit\n",
    "exitFlag_export = 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",
    "print(\"Total time elapsed:\",time.time()-timer,\"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, remove empty entries and dump it all into one **.h5** file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove empty entries when training with a single animal\n",
    "out_df.drop(out_df.index[remove_empty_entries], inplace=True)\n",
    "# reset the indices of the updated dataframe\n",
    "out_df.reset_index(drop=True, inplace=True)\n",
    "out_df.to_hdf(os.path.join(target_dir, \"Data_3D_Pose.hdf5\"),\"df_with_missing\",mode=\"w\")\n",
    "\n",
    "out_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# and finally, dump all the labels into a lookup table\n",
    "with open(os.path.join(target_dir,'label_names.txt'), 'w') as f:\n",
    "    # skip first for, and use only every 5th element\n",
    "    for label in labels:\n",
    "        f.write(label.split(\".\")[0]+\"\\n\")\n",
    "\n",
    "print(\"All done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now that the dataset is generated, let's run a few checks to see if the data contains what we need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "def set_axes_equal(ax):\n",
    "    '''Make axes of 3D plot have equal scale so that spheres appear as spheres,\n",
    "    cubes as cubes, etc..  This is one possible solution to Matplotlib's\n",
    "    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.\n",
    "    \n",
    "    Function based on https://stackoverflow.com/questions/13685386/\n",
    "    matplotlib-equal-unit-length-with-equal-aspect-ratio-z-axis-is-not-equal-to\n",
    "\n",
    "    Input\n",
    "      ax: a matplotlib axis, e.g., as output from plt.gca().\n",
    "    '''\n",
    "\n",
    "    x_limits = ax.get_xlim3d()\n",
    "    y_limits = ax.get_ylim3d()\n",
    "    z_limits = ax.get_zlim3d()\n",
    "\n",
    "    x_range = abs(x_limits[1] - x_limits[0])\n",
    "    x_middle = np.mean(x_limits)\n",
    "    y_range = abs(y_limits[1] - y_limits[0])\n",
    "    y_middle = np.mean(y_limits)\n",
    "    z_range = abs(z_limits[1] - z_limits[0])\n",
    "    z_middle = np.mean(z_limits)\n",
    "\n",
    "    # The plot bounding box is a sphere in the sense of the infinity\n",
    "    # norm, hence I call half the max range the plot radius.\n",
    "    plot_radius = 0.5*max([x_range, y_range, z_range])\n",
    "\n",
    "    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])\n",
    "    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])\n",
    "    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start off by creating a 3D scatter plot of an example sample. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "# to open plot externally and use it interactively\n",
    "show_entry = 0\n",
    "cam_entry = 0\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "\n",
    "display_points_3D = out_df.loc[show_entry][\"key_points_3D\"]\n",
    "display_img = cv2.imread(dataset_img[show_entry])\n",
    "\n",
    "for i,xyz in enumerate(display_points_3D):\n",
    "    if out_df.loc[show_entry][\"visibility\"][i] == 1:\n",
    "        ax.scatter(xyz[0], xyz[1], xyz[2], marker='o',s=10)\n",
    "\n",
    "\"\"\"\n",
    "# also plot the camera location\n",
    "ax.scatter(out_df.loc[cam_entry][\"cam_trans\"][0], \n",
    "           out_df.loc[cam_entry][\"cam_trans\"][1], \n",
    "           out_df.loc[cam_entry][\"cam_trans\"][2], marker='x')\n",
    "\"\"\"\n",
    "\n",
    "ax.set_xlabel('X axis')\n",
    "ax.set_ylabel('Y axis')\n",
    "ax.set_zlabel('Z axis')\n",
    "\n",
    "# use custom function to ensure equal axis proportions\n",
    "set_axes_equal(ax)\n",
    "\n",
    "# opens external plot\n",
    "plt.title(out_df.loc[show_entry][\"file_name\"])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally show that the provided camera intrinsics and extrinsics allow us to project the \n",
    "subject 3D coordinates in world space into the equivalent screen 2D coordinates in pixel space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#%matplotlib qt\n",
    "# file_name\tcam_rot\tcam_trans\tcam_intrinsics\tbounding_box\tkey_points_3D\tkey_points_2D\tvisibility\n",
    "R = np.array(out_df.loc[cam_entry][\"cam_rot\"])\n",
    "T = np.reshape(np.array(out_df.loc[cam_entry][\"cam_trans\"]),(3,1))\n",
    "C = np.array(out_df.loc[cam_entry][\"cam_intrinsics\"])\n",
    "\n",
    "#plt.plot(X_2d[0, :], X_2d[1, :], '.')  # plot the locations of the 3D keypoints in 2D as viewed from the camera\n",
    "#plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "\n",
    "for i, x in enumerate(display_points_3D):\n",
    "    X = np.reshape(np.array(out_df.loc[show_entry][\"key_points_3D\"][i]),(3,-1))\n",
    "\n",
    "    # given the above data, it should be possible to project the 3D points into the corresponding image,\n",
    "    # so they land in the correct position on the image \n",
    "    P = C @ np.hstack([R, T])  # projection matrix\n",
    "    X_hom = np.vstack([X, np.ones(X.shape[1])])  # 3D points in homogenous coordinates\n",
    "\n",
    "    X_hom = P @ X_hom  # project the 3D points\n",
    "    \n",
    "    X_2d = X_hom[:2, :] / X_hom[2, :]  # convert them back to 2D pixel space\n",
    "    \n",
    "    gt_x_2d = out_df.loc[show_entry][\"key_points_2D\"][i][0]\n",
    "    gt_y_2d = out_df.loc[show_entry][\"key_points_2D\"][i][1]\n",
    "    \n",
    "    ax.scatter(gt_x_2d, gt_y_2d, marker='o', s=20)\n",
    "    ax.scatter(X_2d[0], display_img.shape[1]-X_2d[1], marker='^', s=8)\n",
    "\n",
    "ax.set_xlabel('X axis')\n",
    "ax.set_ylabel('Y axis')\n",
    "\n",
    "ax.set_xlim([0,display_img.shape[0]])\n",
    "ax.set_ylim([0,display_img.shape[1]])\n",
    "ax.set_aspect('equal')\n",
    "\n",
    "ax.invert_yaxis()\n",
    "\n",
    "# opens external plot\n",
    "plt.title(str(out_df.loc[show_entry][\"file_name\"]) + \" projected\")\n",
    "plt.show()"
   ]
  },
  {
   "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
}