process_data/convert_to_yolo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "504634c4",
   "metadata": {},
   "source": [
    "# SGG YOLO CONVERTER"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ce5e327",
   "metadata": {},
   "source": [
    "### This script can convert any annotations file in SGG format (.h5) in YOLO format for object detection training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ab662ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import random\n",
    "import h5py\n",
    "import json\n",
    "from collections import Counter\n",
    "from tqdm import tqdm\n",
    "\n",
    "vg_h5 = h5py.File(\"/home/maelic/SGG-Benchmark/datasets/vg/VG-SGG-with-attri.h5\")\n",
    "vg_dict = json.load(open(\"/home/maelic/SGG-Benchmark/datasets/vg/VG-SGG-dicts-with-attri.json\"))\n",
    "image_data = json.load(open('/home/maelic/SGG-Benchmark/datasets/vg/image_data.json'))\n",
    "\n",
    "img_path = '/home/maelic/Documents/Datasets/VG/VG_100K'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3cee7bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_to_idx = vg_dict['idx_to_label']\n",
    "\n",
    "OUT_PATH = \"/home/maelic/Documents/Datasets/VG/VG150/YOLO_anno/\"\n",
    "if not os.path.exists(OUT_PATH):\n",
    "    os.makedirs(OUT_PATH)\n",
    "\n",
    "with open(os.path.join(OUT_PATH, 'classes.txt'), 'w') as f:\n",
    "    for key,v in object_to_idx.items():\n",
    "        f.write(v + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3789b4df",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 108073/108073 [41:21<00:00, 43.55it/s]\n"
     ]
    }
   ],
   "source": [
    "object_to_idx = vg_dict['label_to_idx']\n",
    "dest_folder = os.path.join(OUT_PATH, 'annotations')\n",
    "image_dest_folder = os.path.join(OUT_PATH, 'images_all')\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "if not os.path.exists(dest_folder):\n",
    "    os.makedirs(dest_folder)\n",
    "if not os.path.exists(image_dest_folder):\n",
    "    os.makedirs(image_dest_folder)\n",
    "\n",
    "for idx in tqdm(range(len(vg_h5['img_to_first_box']))):\n",
    "    img_width = image_data[idx]['width']\n",
    "    img_height = image_data[idx]['height']\n",
    "    image_id = image_data[idx]['image_id']\n",
    "    if vg_h5['img_to_first_box'][idx] == -1 or (vg_h5['img_to_last_box'][idx] == vg_h5['img_to_first_box'][idx]):\n",
    "        continue\n",
    "    \n",
    "    for i in range(vg_h5['img_to_first_box'][idx], vg_h5['img_to_last_box'][idx]):\n",
    "\n",
    "        obj_cat = vg_h5['labels'][i][0]-1\n",
    "        assert 0 <= obj_cat < 150\n",
    "        \n",
    "        box = vg_h5['boxes_1024'][i]\n",
    "        # transform coordinates back to original image size\n",
    "        box[0] = box[0] +1\n",
    "        box[1] = box[1] +1\n",
    "        box = box / 1024 * max(img_width, img_height)\n",
    "\n",
    "        # convert to yolo format\n",
    "        x_center = box[0] / img_width\n",
    "        y_center = box[1] / img_height\n",
    "        w = box[2] / img_width\n",
    "        h = box[3] / img_height\n",
    "        \n",
    "        box = [x_center, y_center, w, h]\n",
    "        \n",
    "        # clip to [0, 1]\n",
    "        box = np.clip(box, 0, 1)\n",
    "        \n",
    "        with open(os.path.join(dest_folder, str(image_id) + '.txt'), 'a') as f:\n",
    "            f.write(str(obj_cat) + ' ' + str(box[0]) + ' ' + str(box[1]) + ' ' + str(box[2]) + ' ' + str(box[3])+'\\n')\n",
    "    \n",
    "    # copy image to dest folder\n",
    "    shutil.copyfile(os.path.join(img_path, str(image_id)+'.jpg'), os.path.join(image_dest_folder, str(image_id)+'.jpg'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f2247192",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "base_path = OUT_PATH\n",
    "dest_folder = os.path.join(base_path,'annotations')\n",
    "image_dest_folder = os.path.join(base_path,'images_all')\n",
    "\n",
    "if not os.path.exists(dest_folder):\n",
    "    os.makedirs(dest_folder)\n",
    "\n",
    "if not os.path.exists(image_dest_folder):\n",
    "    os.makedirs(image_dest_folder)\n",
    "\n",
    "images = [os.path.join(image_dest_folder, x) for x in os.listdir(image_dest_folder)]\n",
    "annotations = [os.path.join(dest_folder, x) for x in os.listdir(dest_folder) if x[-3:] == \"txt\"]\n",
    "\n",
    "images.sort()\n",
    "annotations.sort()\n",
    "train_images, val_images, train_annotations, val_annotations = train_test_split(images, annotations, test_size = 0.3, random_state = 1)\n",
    "val_images, test_images, val_annotations, test_annotations = train_test_split(val_images, val_annotations, test_size = 0.5, random_state = 1)\n",
    "\n",
    "root_path = 'images/'\n",
    "folders = [base_path+'/train',base_path+'/test',base_path+'/val']\n",
    "for folder in folders:\n",
    "    os.makedirs(os.path.join(folder,root_path))\n",
    "    \n",
    "root_path = 'labels/'\n",
    "for folder in folders:\n",
    "    os.makedirs(os.path.join(folder,root_path))    \n",
    "\n",
    "def move_files_to_folder(list_of_files, destination_folder):\n",
    "    for f in list_of_files:\n",
    "        try:\n",
    "            shutil.move(f, destination_folder)\n",
    "        except:\n",
    "            print(f)\n",
    "            assert False\n",
    "\n",
    "move_files_to_folder(train_images, base_path+'/train/images/')\n",
    "move_files_to_folder(val_images, base_path+'/val/images/')\n",
    "move_files_to_folder(test_images, base_path+'/test/images/')\n",
    "move_files_to_folder(train_annotations, base_path+'/train/labels/')\n",
    "move_files_to_folder(val_annotations, base_path+'/val/labels/')\n",
    "move_files_to_folder(test_annotations, base_path+'/test/labels/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8e2aeb0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the yaml file\n",
    "classes = vg_dict['idx_to_label'].values()\n",
    "\n",
    "train_path = base_path+'train/'\n",
    "val_path = base_path+'val/'\n",
    "test_path = base_path+'test/'\n",
    "\n",
    "n_classes = len(classes)\n",
    "\n",
    "with open(base_path+'/train.yaml', 'w') as f:\n",
    "    f.write('train: '+train_path+'\\n')\n",
    "    f.write('val: '+val_path+'\\n')\n",
    "    f.write('test: '+test_path+'\\n')\n",
    "    f.write('nc: '+str(n_classes)+'\\n')\n",
    "    f.write('names: '+str(list(classes)))"
   ]
  }
 ],
 "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.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}