UE5_parsers/Generate_YOLO_Style_Dataset_JSON.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# YOLO style dataset parser\n",
    "\n",
    "### The generated output includes the following annotation data:\n",
    "* Classes\n",
    "* Bounding Boxes\n",
    "\n",
    "### Example application(s) (as demonstrated in Plum et al. 2023):\n",
    "* [YOLO v4](https://github.com/AlexeyAB/darknet) _(multi-animal detection and tracking)_\n",
    "\n",
    "### Output structure:\n",
    "* target_dir\n",
    "    * data\n",
    "      * obj\n",
    "        * images\n",
    "        * annotation files\n",
    "      * ###-train.txt _(contains list of training samples)_\n",
    "      * ###-test.txt _(contains list of validation samples)_\n",
    "      * obj.data _(relative dataset path info)_\n",
    "      * obj.names _(all class names)\n",
    "    \n",
    "### Notes:\n",
    "\n",
    "* While we used [YOLO v4](https://github.com/AlexeyAB/darknet) in all our examples, the dataset structure is generally applicable to most YOLO detectors (excluding those that additionally infer pose e.g. YOLO v7)\n",
    "* We also demonstrate how trained models can be used for real-time multi animal tracking applications in our [Blender Multi-Animal Tracking and Pose Estimation Addon - **OmniTrax**](https://github.com/FabianPlum/OmniTrax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib as plt\n",
    "import pathlib\n",
    "\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "\n",
    "import json\n",
    "import threading\n",
    "import queue\n",
    "import sys\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "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-multi\"\n",
    "target_dir = \"../example_data/YOLO\"\n",
    "\n",
    "# Specify which labels to ignore. By default, all keypoints are used to recompute bounding boxes.\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']\n",
    "\n",
    "# The following options are only relevant for performing cross-validation with specific naming conventions.\n",
    "# if the dataset_name is left unspecified (= None), the name will be derived from the original BatchData file\n",
    "dataset_name = \"eq\"\n",
    "dataset_img_name = \"Atta-hist-eq\""
   ]
  },
  {
   "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",
    "enforce_single_class = False # overwrites multiple classes and groups all instances as one\n",
    "\n",
    "cross_validation_split = [5,0]\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",
    "# only use detections where at least 5% of the bounding box is occupied by the respective subject\n",
    "visibility_threshold = 0.05\n",
    "\n",
    "# we can additionally plot the points in the data files to check joint locations\n",
    "plot_joints = False\n",
    "\n",
    "# remember to refine an export folder when saving out your dataset\n",
    "generate_dataset = True\n",
    "\n",
    "# we can enforce the bounding box to centre on the individual instead of being influenced by its orientation\n",
    "# As the groundtruth in real recordings is annotated in the same way this should boost the average accuracy\n",
    "# When enforcing a centred bounding box, list between (on which) which keypoints(s) the centre is supposed to be placed.\n",
    "# In the default insect rig, the centre would lie between the thorax b_t and abdomen b_a_1 keypoint\n",
    "enforce_centred_bboxes = False\n",
    "enforced_centre = [\"b_t\", \"b_a_1\"]\n",
    "\n",
    "# alternatively, we can draw tighter bounding boxes without enforced centres, based on 2D keypoints listed in \"labels\"\n",
    "enforce_tight_bboxes = False # centred OR tight. This option will overwrite \"enforce_centred\" if True"
   ]
  },
  {
   "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",
    "\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",
    "if not enforce_single_class:\n",
    "    # get provided classes to create a dictionary of class IDs and class names\n",
    "    all_classes = []\n",
    "    for subject in colony[\"Subject Variations\"]:\n",
    "        all_classes.append(colony[\"Subject Variations\"][subject][\"Class\"])\n",
    "        \n",
    "    subject_class_names = []\n",
    "    for class_name in all_classes:\n",
    "        # check if exists in unique_list or not and replace any spaces with underscores\n",
    "        class_name = class_name.replace(\" \", \"_\")\n",
    "        if class_name not in subject_class_names:\n",
    "            # append unique classes \n",
    "            subject_class_names.append(class_name)\n",
    "            \n",
    "    # sort classes for repeatability between datasets\n",
    "    subject_class_names.sort()\n",
    "        \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",
    "    subject_classes = dict(sorted(subject_classes.items()))\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\\n\",subject_classes)\n",
    "\n",
    "print(\"\\nLoaded colony file with seed\", colony['Seed']) #,\"and\",len(colony['ID']),\"individuals.\")"
   ]
  },
  {
   "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 simply make this a list of lists as the number of individuals.\n",
    "\n",
    "We will therefore access \"data\" as [frame] [individual] [attribute], where attributes will include [ID,bbox_x_0,bbox_y_0,...]\n",
    "\n",
    "In this instance, we use the bounding box of each individual to first determine its visibility through its occpuancy in the ID pass and second produce a cropped sample for every individual, including its keypoints in the cropped ROI.\n",
    "\n",
    "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.\n",
    "For simplicity we'll assume that at this stage all subjects use the same armature and therefore report the same keypoints.\n",
    "We therefore load the first sample from the list and find the subjects keypoint hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_file = open(dataset_data[0])\n",
    "\n",
    "# returns JSON object as a dictionary\n",
    "sample = json.load(sample_file)\n",
    "sample_file.close()\n",
    "\n",
    "first_entry_key = list(sample[\"iterationData\"][\"subject Data\"][0].keys())[0]\n",
    "labels = list(sample[\"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(\"\\nUsing labels:  \", labels)"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "def fix_bounding_boxes(coords,max_val=[1024,1024],ind_key=None,labels=None):\n",
    "    # fix bounding box coordinates so they do not reach beyond the image\n",
    "    # you can either pass only bounding box coordinates or the entire skeleton coordinates\n",
    "    # The latter will recalculate a tighter bounding box, based on all keypoints\n",
    "    # When recalculating the bounding box based on all keypoints, you can chose to ignore wings.\n",
    "    fixed_coords = []\n",
    "    \n",
    "    if len(coords) == 4:\n",
    "        coords_bbox = coords[:4]\n",
    "    \n",
    "    else:\n",
    "        coords_bbox = [0,0,max_val[0],max_val[1]]\n",
    "        # get all X and Y coordinates to find min and max values for the bounding box\n",
    "        key_x = []\n",
    "        key_y = []\n",
    "        \n",
    "        for key in labels:\n",
    "            key_x.append(coords[ind_key][\"keypoints\"][key][\"2DPos\"][\"x\"])\n",
    "            key_y.append(coords[ind_key][\"keypoints\"][key][\"2DPos\"][\"y\"])\n",
    "        \n",
    "        coords_bbox[0] = max([0,min(key_x)])\n",
    "        coords_bbox[1] = max([0,min(key_y)])\n",
    "        coords_bbox[2] = min([max_val[0],max(key_x)])\n",
    "        coords_bbox[3] = min([max_val[1],max(key_y)])\n",
    "    \n",
    "    for c, coord in enumerate(coords_bbox):\n",
    "        if c == 0 or c == 2:\n",
    "            max_val_temp = max_val[1]\n",
    "        else:\n",
    "            max_val_temp = max_val[0]\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 customThread(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_out = display_img.copy()\n",
    "            \n",
    "            # compute visibility for each individual from ID pass\n",
    "            seg_img = cv2.imread(ID)\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",
    "            if generate_dataset:\n",
    "                img_info = []\n",
    "            \n",
    "            # check if the size of the image and segmentation pass match\n",
    "            if display_img.shape != seg_img.shape:\n",
    "                print(\"Size mismatch of image and segmentation pass for sample\",data_input[1].split(\"/\")[-1],\"!\")\n",
    "            else:\n",
    "                individual_visible = False\n",
    "                \n",
    "                for individual in data[\"iterationData\"][\"subject Data\"]:\n",
    "                    ind_key = list(individual.keys())[0]\n",
    "                    ind_ID = int(ind_key)\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",
    "                    if enforce_tight_bboxes:\n",
    "                        bbox = fix_bounding_boxes(individual, max_val=display_img.shape, ind_key=ind_key, labels=labels)\n",
    "                    else:\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]), np.array([0, 0, ind_ID]))\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(ind_ID), (bbox[0] + 10,bbox[3] - 10), font, fontScale, fontColor, lineType)\n",
    "                    if not enforce_single_class:\n",
    "                        class_ID = subject_classes[colony['Subject Variations'][ind_key][\"Class\"].replace(\" \",\"_\")]\n",
    "                    else:\n",
    "                        # here we use a single class, otherwise this can be replaced by size / scale values\n",
    "                        class_ID = 0\n",
    "                    \n",
    "                    if bbox_occupancy > visibility_threshold:\n",
    "                        \n",
    "                        individual_visible = True\n",
    "                        \n",
    "                        if generate_dataset:\n",
    "                            # now we need to convert the bounding box info into the desired format.\n",
    "                            img_dim = display_img.shape\n",
    "                            \n",
    "                            # [class_ID, centre_x, centre_y, bounding_box_width, bounding_box_height]\n",
    "                            valid_new_x = False\n",
    "                            valid_new_y = False\n",
    "\n",
    "                            if enforce_centred_bboxes:\n",
    "                                # coords of head\n",
    "                                centre_points = []\n",
    "                                for key in enforced_centre:\n",
    "                                    centre_points.append([individual[ind_key][\"keypoints\"][key][\"2DPos\"][\"x\"],\n",
    "                                                         individual[ind_key][\"keypoints\"][key][\"2DPos\"][\"y\"]])\n",
    "                                \n",
    "                                centre_points_arr = np.array(centre_points)\n",
    "                                # compute new centre point\n",
    "                                new_centre_x = np.mean(centre_points_arr[:,0])\n",
    "                                new_centre_y = np.mean(centre_points_arr[:,1])\n",
    "\n",
    "                                if new_centre_x < img_dim[1] and new_centre_x > 0:\n",
    "                                    centre_x = new_centre_x / img_dim[1]\n",
    "                                    valid_new_x = True\n",
    "\n",
    "                                if new_centre_y < img_dim[0] and new_centre_y > 0:\n",
    "                                    centre_y = new_centre_y / img_dim[0]\n",
    "                                    valid_new_y = True\n",
    "\n",
    "                                cv2.circle(display_img, (int(new_centre_x),int(new_centre_y)), \n",
    "                                           radius=3, color=fontColor, thickness=-1)    \n",
    "\n",
    "                            for label in range(int((len(individual)-5)/5)):\n",
    "                                cv2.circle(display_img, (int(individual[label*5+5]),\n",
    "                                                         int(individual[label*5+6])), \n",
    "                                           radius=3, color=fontColor, thickness=-1)    \n",
    "\n",
    "\n",
    "                            bounding_box_width = abs(bbox[2] - bbox[0]) / img_dim[1]\n",
    "                            bounding_box_height = abs(bbox[3] - bbox[1]) / img_dim[0]\n",
    "\n",
    "                            if not valid_new_x or not valid_new_y:\n",
    "                                centre_x = bbox[0] / img_dim[1] + bounding_box_width / 2\n",
    "                                centre_y = bbox[1] / img_dim[0] + bounding_box_height / 2\n",
    "\n",
    "                            img_info.append([class_ID,centre_x,centre_y,bounding_box_width,bounding_box_height])\n",
    "\n",
    "                            cv2.rectangle(display_img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), fontColor, 2)\n",
    "\n",
    "                    else:\n",
    "                        pass\n",
    "                        if len(threadList) == 1: \n",
    "                            print(\"Ah shit, can't see\",ind_ID,class_ID)\n",
    "\n",
    "                    if generate_dataset and individual_visible:\n",
    "\n",
    "                        img_name = target_dir + \"/data/obj/\" + img.split('/')[-1][:-4] + \"_\" + dataset_img_name + \".JPG\"\n",
    "                        cv2.imwrite(img_name, display_img_out)\n",
    "\n",
    "                        with open(target_dir + \"/data/obj/\" + img.split('/')[-1][:-4] + \"_\" + dataset_img_name + \".txt\", \"w\") as f: \n",
    "                            output_txt = []\n",
    "                            if img_info:\n",
    "                                for line in img_info:\n",
    "                                    line_str = ' '.join([str(i) for i in line])\n",
    "                                    output_txt.append(line_str+\"\\n\")\n",
    "                                f.writelines(output_txt)\n",
    "                            else:\n",
    "                                f.write(\"\")\n",
    "\n",
    "                if len(threadList) == 1:\n",
    "                    cv2.imshow(\"labeled image\", cv2.resize(display_img, (int(display_img.shape[1] / 2), \n",
    "                                                                         int(display_img.shape[0] / 2))))\n",
    "                    cv2.waitKey(0)\n",
    "\n",
    "        else:\n",
    "            queueLock.release()\n",
    "            \n",
    "# setup as many threads as there are (virtual) CPUs\n",
    "exitFlag = 0\n",
    "if DEBUG:\n",
    "    threadList = createThreadList(1)\n",
    "else:\n",
    "    threadList = createThreadList(getThreads())\n",
    "print(\"Using\", len(threadList), \"threads to parse data...\")\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_data))\n",
    "threads = []\n",
    "threadID = 1\n",
    "\n",
    "\n",
    "np.random.seed(seed=1)\n",
    "# once colony size can be read from the BatchData file, set the size of ID_colours equal to the colony size\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",
    "if generate_dataset:\n",
    "    from helper.Generate_YOLO_training import createCustomFiles\n",
    "    createCustomFiles(output_folder=target_dir+\"/\",obIDs=subject_class_names, k_fold=cross_validation_split)\n",
    "\n",
    "timer = time.time()\n",
    "\n",
    "# Create new threads\n",
    "for tName in threadList:\n",
    "    thread = customThread(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 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": [
    "Generating cross validation splits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if dataset_name is None:\n",
    "    dataset_name = colony[\"name\"]\n",
    "\n",
    "from helper.Generate_YOLO_training import createCustomFiles\n",
    "for i in range(cross_validation_split[0]):\n",
    "    print(\"Generating split\", i + 1, \"...\")\n",
    "    createCustomFiles(output_folder=target_dir+\"/\",\n",
    "                      obIDs=subject_class_names,\n",
    "                      k_fold=[cross_validation_split[0],i],\n",
    "                      custom_name=dataset_name+str(i+1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, displaying example detections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show example detection\n",
    "ID_colours = np.random.randint(255, size=(255, 3))\n",
    "example_sample_path = \"../example_data/YOLO/data/obj/input-multi_01_example_multi.JPG\"\n",
    "\n",
    "test_img = cv2.imread(example_sample_path)\n",
    "test_img_height,test_img_width = test_img.shape[:-1]\n",
    "\n",
    "with open(example_sample_path[:-3] + \"txt\") as f:\n",
    "    lines = f.readlines()\n",
    "    for l, line in enumerate(lines):\n",
    "        line = line.split(\" \")\n",
    "        # the colours do not correspond to their original IDs as they are not present in this dataset style\n",
    "        # different colours are only used for visualisation purposes to better distinquish between adjacent bounding boxes\n",
    "        ind_color = (int(ID_colours[l,0]),\n",
    "                     int(ID_colours[l,1]),\n",
    "                     int(ID_colours[l,2]))\n",
    "        x,y,w,h = int(float(line[1])*test_img_width),int(float(line[2])*test_img_height),int(float(line[3])*test_img_width),int(float(line[4])*test_img_height)\n",
    "        cv2.rectangle(test_img, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), ind_color, 2)\n",
    "    \n",
    "cv2.imshow(\"test_img\",cv2.resize(test_img,(1024,1024)))\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if generate_dataset:\n",
    "    from helper.Generate_YOLO_training import createCustomFiles\n",
    "    createCustomFiles(output_folder=target_dir+\"/\",obIDs=subject_class_names, k_fold=cross_validation_split)"
   ]
  },
  {
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}