convert_to_yolo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "Author: BinhNVH"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "outputs": [],
   "source": [
    "class_map = {\n",
    "    \"Aortic enlargement\" : 0,\n",
    "    \"Atelectasis\" : 1,\n",
    "    \"Calcification\" : 2,\n",
    "    \"Cardiomegaly\" : 3,\n",
    "    \"Consolidation\" : 4,\n",
    "    \"ILD\" : 5,\n",
    "    \"Infiltration\" : 6,\n",
    "    \"Lung Opacity\" : 7,\n",
    "    \"Nodule/Mass\" : 8,\n",
    "    \"Other lesion\" : 9,\n",
    "    \"Pleural effusion\" : 10,\n",
    "    \"Pleural thickening\" : 11,\n",
    "    \"Pneumothorax\" : 12,\n",
    "    \"Pulmonary fibrosis\" :13\n",
    "}"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "outputs": [],
   "source": [
    "def convert_to_yolo(bbox_file, output_file_yolo):\n",
    "    bbox_df = pd.read_csv(bbox_file)\n",
    "    # Load bounding box data and convert to YOLOv5 format\n",
    "    bbox_df['x_min'] = bbox_df.apply(lambda row: (row.x_min)/row.width, axis =1)\n",
    "    bbox_df['y_min'] = bbox_df.apply(lambda row: (row.y_min)/row.height, axis =1)\n",
    "    bbox_df['x_max'] = bbox_df.apply(lambda row: (row.x_max)/row.width, axis =1)\n",
    "    bbox_df['y_max'] = bbox_df.apply(lambda row: (row.y_max)/row.height, axis =1)\n",
    "    bbox_df['x_mid'] = bbox_df.apply(lambda row: (row.x_max+row.x_min)/2, axis =1)\n",
    "    bbox_df['y_mid'] = bbox_df.apply(lambda row: (row.y_max+row.y_min)/2, axis =1)\n",
    "    bbox_df['w'] = bbox_df.apply(lambda row: (row.x_max-row.x_min), axis =1)\n",
    "    bbox_df['h'] = bbox_df.apply(lambda row: (row.y_max-row.y_min), axis =1)\n",
    "    bbox_df['area'] = bbox_df['w']*bbox_df['h']\n",
    "    new_bbox_df = bbox_df[['class_id', 'x_mid', 'y_mid', 'w', 'h']]\n",
    "\n",
    "    # Group bounding boxes by image ID and save as separate txt files\n",
    "    image_ids = bbox_df[\"image_id\"].unique()\n",
    "    for image_id in image_ids:\n",
    "        output_file = os.path.join(output_file_yolo, f\"{image_id}.txt\")\n",
    "        image_bboxes = new_bbox_df[bbox_df[\"image_id\"] == image_id]\n",
    "        image_bboxes.to_csv(output_file, header=None, index=None, sep=\" \")\n",
    "    print(\"Conversion complete!\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "outputs": [],
   "source": [
    "os.makedirs(r'E:\\XR-AI\\datasets\\vinbigdata\\yolo_label',exist_ok = True)\n",
    "input = r'E:\\XR-AI\\datasets\\vinbigdata\\train.csv'\n",
    "output = r'E:\\XR-AI\\datasets\\vinbigdata\\yolo_label'\n",
    "convert_to_yolo(input,output)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "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
}