[ea94d1]: / FracAtlas / Utilities / yolo2voc.ipynb

Download this file

86 lines (85 with data), 3.4 kB

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "# Define the path of YOLO folder and PASCAL VOC folder\n",
    "yolo_path = 'YOLO'\n",
    "pascal_path = 'PASCAL VOC'\n",
    "\n",
    "# Read the classes from the classes.txt file\n",
    "with open(os.path.join(yolo_path, 'classes.txt'), 'r') as f:\n",
    "    classes = f.read().splitlines()\n",
    "\n",
    "# Create the PASCAL VOC folder structure\n",
    "if not os.path.exists(pascal_path):\n",
    "    os.makedirs(os.path.join(pascal_path, 'Annotations'))\n",
    "\n",
    "# Convert each YOLO annotation file to PASCAL VOC format\n",
    "for file_name in os.listdir(os.path.join(yolo_path, 'Annotations')):\n",
    "    # Read the YOLO annotation file\n",
    "    with open(os.path.join(yolo_path, 'Annotations', file_name), 'r') as f:\n",
    "        lines = f.read().splitlines()\n",
    "\n",
    "    # Create the PASCAL VOC annotation file\n",
    "    annotation = ET.Element('annotation')\n",
    "    ET.SubElement(annotation, 'folder').text = 'JPEGImages'\n",
    "    ET.SubElement(annotation, 'filename').text = file_name[:-4] + '.jpg'\n",
    "    ET.SubElement(annotation, 'source').text = 'unknown'\n",
    "    size = ET.SubElement(annotation, 'size')\n",
    "    ET.SubElement(size, 'width').text = '0'\n",
    "    ET.SubElement(size, 'height').text = '0'\n",
    "    ET.SubElement(size, 'depth').text = '3'\n",
    "    ET.SubElement(annotation, 'segmented').text = '0'\n",
    "\n",
    "    # Add object annotation for each object in the YOLO annotation file\n",
    "    for line in lines:\n",
    "        class_id, x_center, y_center, width, height = map(float, line.split())\n",
    "        class_name = classes[int(class_id)]\n",
    "        object_annotation = ET.SubElement(annotation, 'object')\n",
    "        ET.SubElement(object_annotation, 'name').text = class_name\n",
    "        ET.SubElement(object_annotation, 'pose').text = 'Unspecified'\n",
    "        ET.SubElement(object_annotation, 'truncated').text = '0'\n",
    "        ET.SubElement(object_annotation, 'difficult').text = '0'\n",
    "        bndbox = ET.SubElement(object_annotation, 'bndbox')\n",
    "        xmin = x_center - width/2\n",
    "        ymin = y_center - height/2\n",
    "        xmax = x_center + width/2\n",
    "        ymax = y_center + height/2\n",
    "        ET.SubElement(bndbox, 'xmin').text = str(int(xmin))\n",
    "        ET.SubElement(bndbox, 'ymin').text = str(int(ymin))\n",
    "        ET.SubElement(bndbox, 'xmax').text = str(int(xmax))\n",
    "        ET.SubElement(bndbox, 'ymax').text = str(int(ymax))\n",
    "\n",
    "    # Save the PASCAL VOC annotation file\n",
    "    tree = ET.ElementTree(annotation)\n",
    "    tree.write(os.path.join(pascal_path, 'Annotations', file_name[:-4] + '.xml'))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "77705bf0baf5524fde679a6aa5ebe94fde81fd480e2e327e6e007bf7804e3f54"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}