dataset2yolo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from imutils import paths \n",
    "from shutil import copyfile\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import cv2\n",
    "import csv\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_dir(dirName):\n",
    "    # Create a target directory & all intermediate \n",
    "    # directories if they don't exists\n",
    "    if not os.path.exists(dirName):\n",
    "        os.makedirs(dirName, exist_ok = True)\n",
    "        print(\"[INFO] Directory \" , dirName ,  \" created\")\n",
    "    else:\n",
    "        print(\"[INFO] Directory \" , dirName ,  \" already exists\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOURCE_PATH = 'TFW/'\n",
    "TARGET_PATH = 'dataset/YOLOv5'\n",
    "\n",
    "TRAIN_IMG_DIR = 'train/images'\n",
    "TRAIN_LAB_DIR = 'train/labels'\n",
    "make_dir(os.path.join(TARGET_PATH, TRAIN_IMG_DIR))\n",
    "make_dir(os.path.join(TARGET_PATH, TRAIN_LAB_DIR))\n",
    "\n",
    "TEST_IMG_DIR = 'test/images'\n",
    "TEST_LAB_DIR = 'test/labels'\n",
    "make_dir(os.path.join(TARGET_PATH, TEST_IMG_DIR))\n",
    "make_dir(os.path.join(TARGET_PATH, TEST_LAB_DIR))\n",
    "\n",
    "VAL_IMG_DIR = 'val/images'\n",
    "VAL_LAB_DIR = 'val/labels'\n",
    "make_dir(os.path.join(TARGET_PATH, VAL_IMG_DIR))\n",
    "make_dir(os.path.join(TARGET_PATH, VAL_LAB_DIR))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert(source_path, target_path, set_='train', env='indoor'):\n",
    "    img_path = os.path.join(source_path, set_, env+\"_gray\")\n",
    "    csv_file = os.path.join(source_path, set_, '{}.csv'.format(env))\n",
    "    \n",
    "    # import labels as a numpy array\n",
    "    labels = pd.read_csv(csv_file,header=0).to_numpy()\n",
    "\n",
    "    # initialize iterators\n",
    "    iter1 = 0 \n",
    "    iter2 = 0\n",
    "\n",
    "    # init face counter\n",
    "    face_count = 0\n",
    "\n",
    "    # loop over the rows\n",
    "    while iter1 < len(labels):\n",
    "        # load the image\n",
    "        imageName = labels[iter1][0]\n",
    "        \n",
    "        source_img = os.path.join(img_path, imageName)\n",
    "        image = cv2.imread(source_img)\n",
    "        H, W = image.shape[:2]\n",
    "        \n",
    "        target_img = os.path.join(target_path, set_, 'images', imageName)\n",
    "        target_txt = open(os.path.join(target_path, set_,\n",
    "                                'labels/{}.txt'.format(imageName.split('.png')[0])), \"w\")\n",
    "        \n",
    "        copyfile(source_img, target_img)\n",
    "        \n",
    "        # loop over the labels of this image\n",
    "        while labels[iter1][0] == labels[iter2][0]:\n",
    "            # extract coordinates of the bounding box and five facial landmarks\n",
    "            xs, ys, xe, ye, p1x, p1y, p2x, p2y, p3x, p3y, p4x, p4y, p5x, p5y = labels[iter2][1:]\n",
    "            label = 'face {} {} {} {} {} {} {} {} {} {} {} {} {} {}\\n'.format(xs, ys, xe-xs, ye-ys, p1x, p1y, p2x, p2y, \n",
    "                                                                         p3x, p3y, p4x, p4y, p5x, p5y)\n",
    "            w = xe - xs\n",
    "            h = ye - ys\n",
    "                \n",
    "            xc = xs + w / 2\n",
    "            yc = ys + h / 2\n",
    "                \n",
    "            label = '0 {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f}\\n'.format(\n",
    "                    xc/W, yc/H, w/W, h/H, p1x/W, p1y/H, p2x/W, p2y/H, p3x/W, p3y/H, p4x/W, p4y/H, p5x/W, p5y/H)\n",
    "                    \n",
    "            target_txt.write(label)\n",
    "            face_count += 1\n",
    "                \n",
    "            iter2 += 1\n",
    "                \n",
    "            # break the loop if we exceeded\n",
    "            # the last row\n",
    "            if iter2 == len(labels):\n",
    "                break\n",
    "        \n",
    "        # equalize iterators\n",
    "        iter1 = iter2\n",
    "\n",
    "    print(iter2, face_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='train', env='indoor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='train', env='outdoor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='val', env='indoor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='val', env='outdoor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='test', env='indoor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SOURCE_PATH, TARGET_PATH, set_='test', env='outdoor')"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}