{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "### YOLO to Standard format conversion script"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "base_dir = os.getcwd()\n",
    "\n",
    "# Paths to the original dataset\n",
    "ORIGINAL_TRAIN_DIR = os.path.join(base_dir,'dataset/train/images')\n",
    "ORIGINAL_VAL_DIR = os.path.join(base_dir,'dataset/valid/images')\n",
    "ORIGINAL_TEST_DIR = os.path.join(base_dir,'dataset/test/images')\n",
    "TRAIN_LABELS_DIR = os.path.join(base_dir,'dataset/train/labels')\n",
    "VAL_LABELS_DIR = os.path.join(base_dir,'dataset/valid/labels')\n",
    "TEST_LABELS_DIR = os.path.join(base_dir,'dataset/test/labels')\n",
    "\n",
    "# Paths to the cropped images based on labels\n",
    "CROPPED_TRAIN_DIR = os.path.join(base_dir,'cropped_dataset/train')\n",
    "CROPPED_VAL_DIR = os.path.join(base_dir,'cropped_dataset/valid')\n",
    "CROPPED_TEST_DIR = os.path.join(base_dir,'cropped_dataset/test')\n",
    "\n",
    "def preprocess_dataset(images_dir, labels_dir, cropped_dir):\n",
    "    if not os.path.exists(cropped_dir):\n",
    "        os.makedirs(cropped_dir)\n",
    "\n",
    "    for label in ['awake', 'sleepy']:\n",
    "        label_dir = os.path.join(cropped_dir, label)\n",
    "        if not os.path.exists(label_dir):\n",
    "            os.makedirs(label_dir)\n",
    "\n",
    "    for img_name in tqdm(os.listdir(images_dir), desc=f'Processing images in {images_dir}'):\n",
    "        img_path = os.path.join(images_dir, img_name)\n",
    "        label_path = os.path.join(labels_dir, os.path.splitext(img_name)[0] + '.txt')\n",
    "\n",
    "        if not os.path.exists(label_path):\n",
    "            continue\n",
    "\n",
    "        with open(label_path, 'r') as f:\n",
    "            labels = f.readlines()\n",
    "\n",
    "        image = cv2.imread(img_path)\n",
    "        height, width, _ = image.shape\n",
    "\n",
    "        for label in labels:\n",
    "            if label.strip() == '':\n",
    "                continue\n",
    "\n",
    "            label_parts = label.strip().split()\n",
    "            if len(label_parts) != 5:\n",
    "                continue\n",
    "\n",
    "            class_id, x_center, y_center, bbox_width, bbox_height = map(float, label_parts)\n",
    "\n",
    "            # Convert to pixel coordinates\n",
    "            x1 = int((x_center - bbox_width / 2) * width)\n",
    "            y1 = int((y_center - bbox_height / 2) * height)\n",
    "            x2 = int((x_center + bbox_width / 2) * width)\n",
    "            y2 = int((y_center + bbox_height / 2) * height)\n",
    "\n",
    "            # Ensure bounding box coordinates are within image boundaries\n",
    "            x1 = max(0, x1)\n",
    "            y1 = max(0, y1)\n",
    "            x2 = min(width, x2)\n",
    "            y2 = min(height, y2)\n",
    "\n",
    "            # Crop the face region\n",
    "            face_img = image[y1:y2, x1:x2]\n",
    "\n",
    "            # Determine the label for saving the cropped face\n",
    "            if int(class_id) == 0:\n",
    "                face_label_dir = os.path.join(cropped_dir, 'awake')\n",
    "            else:\n",
    "                face_label_dir = os.path.join(cropped_dir, 'sleepy')\n",
    "\n",
    "            # Save the cropped face image\n",
    "            face_img_path = os.path.join(face_label_dir, img_name)\n",
    "            cv2.imwrite(face_img_path, face_img)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Preprocess the datasets\n",
    "preprocess_dataset(ORIGINAL_TRAIN_DIR, TRAIN_LABELS_DIR, CROPPED_TRAIN_DIR)\n",
    "preprocess_dataset(ORIGINAL_VAL_DIR, VAL_LABELS_DIR, CROPPED_VAL_DIR)\n",
    "preprocess_dataset(ORIGINAL_TEST_DIR, TEST_LABELS_DIR, CROPPED_TEST_DIR)\n"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}