PubTables-1M/voc2coco_structure.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "261b8e28-12b2-4a43-af71-35310b03b68e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install xmltodict  # -i https://opentuna.cn/pypi/web/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bff9744-9c68-4d0e-ab53-dc50bb7900fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import xmltodict\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ef97804-0902-4846-b662-0cc0d4949885",
   "metadata": {},
   "outputs": [],
   "source": [
    "voc_base_dir = 'PubTables1M-Structure-PASCAL-VOC'\n",
    "coco_base_dir = 'PubTables1M-Structure-COCO'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "329ed238-7340-466c-8336-0d52e22338b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $coco_base_dir\n",
    "!mkdir -p $coco_base_dir/images/train\n",
    "!mkdir -p $coco_base_dir/images/val\n",
    "!mkdir -p $coco_base_dir/images/test\n",
    "!mkdir -p $coco_base_dir/labels/train\n",
    "!mkdir -p $coco_base_dir/labels/val\n",
    "!mkdir -p $coco_base_dir/labels/test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b44964-6d11-4906-b00c-0de6c806b56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "max_processes = multiprocessing.cpu_count()\n",
    "print('max_processes:', max_processes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e18c4548-3ed8-4e4e-ab44-6b593a6af331",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names_map = {'table':0, 'table column':1, 'table row':2, 'table column header':3, 'table projected row header':4, 'table spanning cell':5, 'table grid cell':6}\n",
    "\n",
    "def voc2coco_single(xml_label_filename):\n",
    "    xml_label_filename = xml_label_filename.strip()\n",
    "    if xml_label_filename.endswith('xml'):\n",
    "        file_object = open(os.path.join(voc_base_dir, xml_label_filename), encoding='utf-8')                                                                                                            \n",
    "        try:\n",
    "            all_the_xmlStr = file_object.read()\n",
    "        finally:\n",
    "            file_object.close()\n",
    "        convertedDict = xmltodict.parse(all_the_xmlStr)\n",
    "        # print(convertedDict)\n",
    "        # print(len(convertedDict['annotation']['object']))\n",
    "        fix_width = int(convertedDict['annotation']['size']['width'])\n",
    "        fix_height = int(convertedDict['annotation']['size']['height'])\n",
    "        if 'object' in convertedDict['annotation']:\n",
    "            objs = convertedDict['annotation']['object']\n",
    "            if not isinstance(objs,list):\n",
    "                objs = [objs]\n",
    "#                 print('objs:', objs)\n",
    "            with open(os.path.join(coco_base_dir, 'labels', xml_label_filename[:-4]+'.txt'), 'w') as fout:\n",
    "                for annotation in objs:\n",
    "                    if annotation['name'] not in class_names_map:\n",
    "                        class_names_map[annotation['name']] = len(class_names_map)\n",
    "                    class_id = class_names_map[annotation['name']]\n",
    "                    # class_id = 0\n",
    "\n",
    "                    xmin = int(float(annotation['bndbox']['xmin']))\n",
    "                    ymin = int(float(annotation['bndbox']['ymin']))\n",
    "                    xmax = int(float(annotation['bndbox']['xmax']))\n",
    "                    ymax = int(float(annotation['bndbox']['ymax']))\n",
    "\n",
    "                    w = xmax-xmin\n",
    "                    h = ymax-ymin\n",
    "\n",
    "                    if w>0 and h>0:\n",
    "                        center_x = (xmin+xmax)/2\n",
    "                        center_y = (ymin+ymax)/2\n",
    "                        fout.write(str(class_id)+' '+str(center_x/fix_width)+' '+str(center_y/fix_height)+' '+str(w/fix_width)+' '+str(h/fix_height)+'\\n')\n",
    "    else:\n",
    "        print('[BUG] xml_label_filename:', xml_label_filename)\n",
    "\n",
    "def voc2coco(split='train'):\n",
    "    with open(os.path.join(voc_base_dir, split+'_filelist.txt'), 'r') as fin:\n",
    "        xml_label_filenames = fin.readlines()\n",
    "    print(split, len(xml_label_filenames))\n",
    "    Parallel(n_jobs=max_processes)(delayed(voc2coco_single)(xml_label_filename) for xml_label_filename in tqdm(xml_label_filenames))        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17b6ad73-7d4b-41c1-b82c-311c391bb3b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "voc2coco('train')\n",
    "voc2coco('val')\n",
    "voc2coco('test')\n",
    "\n",
    "print('class_names_map:', class_names_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26cb6ad5-9c94-4eba-a6f6-02a31881e9ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_images(split='train'):\n",
    "    with open(os.path.join(voc_base_dir, split+'_filelist.txt'), 'r') as fin:\n",
    "        xml_label_filenames = fin.readlines()\n",
    "    print(split, len(xml_label_filenames))\n",
    "    for xml_label_filename in tqdm(xml_label_filenames):\n",
    "        xml_label_filename = xml_label_filename.strip()\n",
    "        if xml_label_filename.endswith('xml'):\n",
    "            jpg_filename = os.path.join(voc_base_dir, 'images', xml_label_filename.split('/')[1].replace('xml', 'jpg'))\n",
    "            new_jpg_filename = os.path.join(coco_base_dir, 'images', xml_label_filename.replace('xml', 'jpg'))\n",
    "            # print(jpg_filename, new_jpg_filename)\n",
    "            # shutil.copy(jpg_filename, new_jpg_filename)\n",
    "            shutil.move(jpg_filename, new_jpg_filename)\n",
    "            # break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cafe5e3-ef2d-4923-9df5-d2f91bcf548a",
   "metadata": {},
   "outputs": [],
   "source": [
    "copy_images('train')\n",
    "copy_images('val')\n",
    "copy_images('test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e7909db-8ec6-4e0b-aa33-7ac57ab79b77",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}