UE5_parsers/Generate_mmlab_segmentation_Style_Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MMLAB / MMSEG style dataset parser\n",
    "\n",
    "### The generated output includes the following annotation data:\n",
    "* Classes\n",
    "* Pixel masks\n",
    "\n",
    "### Example application(s) (as demonstrated in Plum et al. 2023):\n",
    "* PSPNet _(semantic segmentation)_\n",
    "* UperNet + SWIN Transformers _(semantic segmentation)_\n",
    "\n",
    "### Output structure:\n",
    "* target_dir\n",
    "  * **images**  _(containing a copy of the original input images)_\n",
    "  * **annotations**  _(containing the segmentation maps)_\n",
    "  * **imageLists**  _(containing three txt files which contain the names of the images separated into all.txt / test.txt / train.txt_\n",
    "    \n",
    "### Notes:\n",
    "\n",
    "* In our examples, we use the [mmsegmentation](https://github.com/open-mmlab/mmsegmentation) repo by [OpenMMLAB](https://openmmlab.com/)\n",
    "* This parser in particular has been used to train models for binary segmentation problems, grouping images into background (0) and specimen (1), but should generally be applicable to arbitrary numbers of classes (<255)\n",
    "* as the colour value of the pixel masks correspond to the respective classes (0 to 255), especially binary images may appear entirely black but will contain all required information"
   ]
  },
  {
   "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",
    "import pathlib\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib as plt\n",
    "\n",
    "from PIL import Image\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.model_selection import KFold\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-multi\"\n",
    "target_dir = \"../example_data/MMSEG\""
   ]
  },
  {
   "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",
    "# number of pixels added to pad masks to avoid cutting off contours\n",
    "# in our examples, we use a padding of 5 pixels during training and erode the resulting masks by 5 pixels during inference\n",
    "mask_padding = 0\n",
    "\n",
    "enforce_single_class = False # overwrites multiple classes and groups all instances as one\n",
    "\n",
    "# determine the proportion of a bounding box that needs to be filled before considering the visibility as too low\n",
    "visibility_threshold = 0.1\n",
    "\n",
    "# test split (what split is witheld)\n",
    "amount_test = 0.1"
   ]
  },
  {
   "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": {},
   "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[-1].split(\".\")[0] == \"ID\":\n",
    "        dataset_ID.append(loc)\n",
    "    elif file_info[-1].split(\".\")[0]  == \"depth\":\n",
    "        dataset_depth.append(loc)\n",
    "    elif file_info[-1].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",
    "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",
    "    subject_classes = {}\n",
    "    for id,sbj in enumerate(subject_class_names):\n",
    "        subject_classes[str(sbj)] = id\n",
    "else:\n",
    "    subject_class_names = [\"insect_0\"] #np.array([int(0)], dtype=int)\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": [
    "Create required output folders used in the mmlab segmentation conventions\n",
    "- **/images/** (containing a copy of the original input images)\n",
    "- **/annotations/** (containing the segmentation maps)\n",
    "- **/imageLists/**(containing three txt files which contain the names of the images separated into all.txt / test.txt / train.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create output folders\n",
    "output_folders = [\"images\", \"annotations\", \"imageLists\"]\n",
    "for i, f in enumerate(output_folders):\n",
    "    if not os.path.exists(target_dir + \"/\" + f):\n",
    "        os.mkdir(target_dir + \"/\" + f)\n",
    "        \n",
    "\"\"\"\n",
    "    if i < 2:\n",
    "        if not os.path.exists(target_dir + \"/\" + f + \"/train\"):\n",
    "            os.mkdir(target_dir + \"/\" + f + \"/train\")\n",
    "        if not os.path.exists(target_dir + \"/\" + f + \"/test\"):\n",
    "            os.mkdir(target_dir + \"/\" + f + \"/test\")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a test and a train image list, according to the specified split.\n",
    "\n",
    "When writing out the images, automatically place them into their respective folders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageLists_all = target_dir + \"/imageLists/all.txt\"\n",
    "imageLists_train = target_dir + \"/imageLists/train.txt\"\n",
    "imageLists_test = target_dir + \"/imageLists/test.txt\"\n",
    "\n",
    "imageLists_all_orig = [img.split('/')[-1][:-4] + \"_synth\" for img in dataset_img]\n",
    "\n",
    "imageLists_all_orig_shuffle = shuffle(imageLists_all_orig, random_state=0)\n",
    "num_train_examples = int(np.floor(len(imageLists_all_orig_shuffle) * (1 - amount_test)))\n",
    "\n",
    "print(\"Using\", num_train_examples, \"training images and\",\n",
    "      int(np.floor(len(imageLists_all_orig_shuffle) - (len(imageLists_all_orig_shuffle) * (1 - amount_test)))), \"test images. (\" + str(amount_test * 100),\n",
    "      \"%)\")\n",
    "\n",
    "files_train = imageLists_all_orig_shuffle[0:num_train_examples]\n",
    "files_test = imageLists_all_orig_shuffle[num_train_examples:]\n",
    "\n",
    "with open(imageLists_all, 'w') as f:\n",
    "    for line in imageLists_all_orig_shuffle:\n",
    "        f.write(line)\n",
    "        f.write('\\n')\n",
    "        \n",
    "with open(imageLists_train, 'w') as f:\n",
    "    for line in files_train:\n",
    "        f.write(line)\n",
    "        f.write('\\n')\n",
    "        \n",
    "with open(imageLists_test, 'w') as f:\n",
    "    for line in files_test:\n",
    "        f.write(line)\n",
    "        f.write('\\n')"
   ]
  },
  {
   "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": [
    "# 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 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",
    "            output_mask = np.zeros(img_shape[:2], dtype=np.uint8)\n",
    "            \n",
    "            # only add images that contain visibile individuals\n",
    "            is_empty = True\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",
    "                for im, individual in enumerate(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",
    "\n",
    "                    contours_lowpoly = []\n",
    "\n",
    "                    #try:\n",
    "                    ID_mask = cv2.inRange(seg_img, np.array([0, 0, ind_ID - 0]), np.array([0, 0, ind_ID + 0]))\n",
    "                    indivual_occupancy = cv2.countNonZero(ID_mask)\n",
    "\n",
    "                    # the kernel size for both dilation and median blur are to be determined by the bbounding boxes relative size\n",
    "                    rel_size = ((bbox[2] - bbox[0]) / display_img.shape[0] + (bbox[3] - bbox[1]) / display_img.shape[0]) / 2\n",
    "                    # values range from 0 (tiny) to 1 (huge)\n",
    "                    # required smoothing 5 to 95\n",
    "                    rel_size_root = int(round((5 * rel_size)/2.)*2 + 1) # round to next odd integer\n",
    "                    #print(\"img:\", i, \"individual:\", im, \"rel_size\", rel_size, rel_size_root)\n",
    "\n",
    "                    # to simplify the generated masks and counter compression artifacts the original mask is dilated\n",
    "                    # https://docs.opencv.org/3.4/db/df6/tutorial_erosion_dilatation.html\n",
    "                    kernel = np.ones((rel_size_root, rel_size_root), 'uint8')\n",
    "                    ID_mask_dilated = cv2.dilate(ID_mask, kernel, iterations=1)\n",
    "                    # use median blur to further smooth the edges of the binary mask\n",
    "                    ID_mask_dilated = cv2.medianBlur(ID_mask_dilated,rel_size_root)\n",
    "\n",
    "                    # pad segmentation subwindow to prevent contours from being cut off\n",
    "\n",
    "                    if mask_padding != 0:\n",
    "                        ID_mask_dilated_padded = np.zeros([ID_mask_dilated.shape[0] + mask_padding * 2 , ID_mask_dilated.shape[1] + mask_padding * 2], 'uint8')\n",
    "                        ID_mask_dilated_padded[mask_padding:-mask_padding,mask_padding:-mask_padding] = ID_mask_dilated\n",
    "                        ID_mask_dilated = ID_mask_dilated_padded\n",
    "\n",
    "                    \n",
    "\n",
    "\n",
    "                    if DEBUG:\n",
    "                        print(\"\\nindividual\",im,ID_mask_dilated.dtype)\n",
    "                        print(hierarchy)\n",
    "                        # draw the contours on the empty image\n",
    "                        seg_img_display = seg_img.copy()\n",
    "                        cv2.imshow(\"mask: \", ID_mask_dilated)\n",
    "                        cv2.drawContours(seg_img_display[bbox[1]:bbox[3],bbox[0]:bbox[2]], contours, -1, (255,0,0), 3)\n",
    "                        cv2.imshow(\"segmentation: \", seg_img_display[bbox[1]:bbox[3],bbox[0]:bbox[2]])\n",
    "                        cv2.waitKey(0)\n",
    "\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",
    "                    if not enforce_single_class:\n",
    "                        class_ID = subject_classes[colony['Subject Variations'][ind_key][\"Class\"].replace(\" \",\"_\")] + 1\n",
    "                    else:\n",
    "                        # here we use a single class, otherwise this can be replaced by size / scale values\n",
    "                        class_ID = 1 # as the background = 0\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",
    "                        #output_mask += np.array((ID_mask_dilated/255) * class_ID, np.uint8)\n",
    "                        \"\"\"\n",
    "                        segmentation code goes here\n",
    "                        - get mask for every individual\n",
    "                        - set all values in out_mask = class_ID which correspond to the individual\n",
    "                        \"\"\"\n",
    "                        \n",
    "                        output_mask = np.where(ID_mask_dilated > 0, \n",
    "                                               np.array((ID_mask_dilated/255) * (class_ID), np.uint8),\n",
    "                                               output_mask)\n",
    "                        \"\"\"\n",
    "                        output_mask = np.where(ID_mask_dilated > 0, \n",
    "                                               ID_mask_dilated,\n",
    "                                               output_mask)\n",
    "                        \"\"\"\n",
    "                    else:\n",
    "                        pass\n",
    "\n",
    "            # uncomment to show resulting bounding boxes and masks\n",
    "            if DEBUG:\n",
    "                cv2.imshow(\"segmentation: \" ,cv2.resize(seg_img_display, (int(seg_img.shape[1] / 2), \n",
    "                                                                  int(seg_img.shape[0] / 2))))\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(1)\n",
    "            \n",
    "            img_name = img.split('/')[-1][:-4] + \"_synth\"\n",
    "            \n",
    "            \"\"\"\n",
    "            if img_name in files_train:\n",
    "                img_out_path = target_dir + \"/images/train/\" + img_name + \".jpg\"\n",
    "                mask_out_path = target_dir + \"/annotations/train/\" + img_name + \"_labelTrainIds.png\"\n",
    "            else:\n",
    "                img_out_path = target_dir + \"/images/test/\" + img_name + \".jpg\"\n",
    "                mask_out_path = target_dir + \"/annotations/test/\" + img_name + \"_labelTrainIds.png\"\n",
    "            \"\"\"\n",
    "            img_out_path = target_dir + \"/images/\" + img_name + \".jpg\"\n",
    "            mask_out_path = target_dir + \"/annotations/\" + img_name + \"_labelTrainIds.png\"\n",
    "                \n",
    "            cv2.imwrite(img_out_path, display_img)\n",
    "            Image.fromarray(output_mask).save(mask_out_path, 'PNG')\n",
    "            print(\"Saved\", img_name)\n",
    "            \n",
    "        else:\n",
    "            queueLock.release()\n",
    "            \n",
    "# setup as many threads as there are (virtual) CPUs\n",
    "exitFlag_export = 0\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",
    "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 refine an export folder when saving out your dataset\n",
    "generate_dataset = True\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",
    "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 samples\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": "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
}