data2yolo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the necessary packages\n",
    "from imutils import paths\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "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": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATASET_PATH = 'dataset'\n",
    "\n",
    "dirs = os.listdir(DATASET_PATH)\n",
    "\n",
    "len(dirs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['icartoon', 'animalweb', 'tfw', 'widerface']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Directory  dataset/icartoon/train/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/icartoon/val/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/icartoon/test/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/animalweb/train/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/animalweb/val/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/animalweb/test/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/tfw/train/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/tfw/val/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/tfw/test/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/widerface/train/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/widerface/val/labels_yolo  already exists\n",
      "[INFO] Directory  dataset/widerface/test/labels_yolo  already exists\n"
     ]
    }
   ],
   "source": [
    "for d in dirs:\n",
    "    for set_ in ['train', 'val', 'test']:\n",
    "        make_dir(os.path.join(DATASET_PATH, d, set_, 'labels_yolo'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "icartoon train\n",
      "icartoon val\n",
      "icartoon test\n",
      "animalweb train\n",
      "animalweb val\n",
      "animalweb test\n",
      "tfw train\n",
      "tfw val\n",
      "tfw test\n",
      "widerface train\n",
      "widerface val\n",
      "widerface test\n"
     ]
    }
   ],
   "source": [
    "for d in dirs:\n",
    "    for set_ in ['train', 'val', 'test']:\n",
    "        annotPaths = list(paths.list_files(os.path.join(DATASET_PATH, d, set_, 'labels')))\n",
    "        annotPaths = sorted(annotPaths)\n",
    "        \n",
    "        imagePaths = list(paths.list_files(os.path.join(DATASET_PATH, d, set_, 'images')))\n",
    "        imagePaths = sorted(imagePaths)\n",
    "        \n",
    "        print(d, set_)\n",
    "        \n",
    "        for imagePath, annotPath in zip(imagePaths, annotPaths):\n",
    "            image = cv2.imread(imagePath)\n",
    "            H, W = image.shape[:2]\n",
    "            \n",
    "            f = open(annotPath, 'r')\n",
    "            lines = f.readlines()\n",
    "            \n",
    "            yolo_txt_path = os.path.join(DATASET_PATH, d, set_, 'labels_yolo', annotPath.split('/')[-1])\n",
    "            yolo_txt = open(yolo_txt_path, \"w\")\n",
    "            \n",
    "            for line in lines:\n",
    "                line = line.split('\\n')[0]\n",
    "                line = line.split(' ')\n",
    "\n",
    "                coord = [float(i) for i in line[1:]]\n",
    "\n",
    "                xs, ys, w, h, p1x, p1y, p2x, p2y, p3x, p3y, p4x, p4y, p5x, p5y = coord\n",
    "                \n",
    "                xc = xs + w / 2\n",
    "                yc = ys + h / 2\n",
    "                \n",
    "                label_yolo = '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",
    "                yolo_txt.write(label_yolo)\n",
    "                \n",
    "            yolo_txt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in dirs:\n",
    "    for set_ in ['train', 'val', 'test']:\n",
    "        src = os.path.join(DATASET_PATH, d, set_, 'labels')\n",
    "        dst = os.path.join(DATASET_PATH, d, set_, 'labels_eval')\n",
    "        os.rename(src, dst)\n",
    "        \n",
    "        src = os.path.join(DATASET_PATH, d, set_, 'labels_yolo')\n",
    "        dst = os.path.join(DATASET_PATH, d, set_, 'labels')\n",
    "        os.rename(src, dst)"
   ]
  },
  {
   "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
}