Spaces:
Build error
Build error
File size: 4,590 Bytes
19ee737 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | {
"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
}
|