convert_segmentation_ade.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle as pkl\n",
    "from io import BytesIO\n",
    "import numpy as np\n",
    "import base64\n",
    "from PIL import Image, ImageFile\n",
    "from pathlib import Path\n",
    "\n",
    "from tqdm import tqdm\n",
    "from multiprocessing import Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "shifted_id = {}\n",
    "for seg_idx in range(150):\n",
    "    val = seg_idx + 1\n",
    "    shifted_id[seg_idx] = val\n",
    "shifted_id[150] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = 'validation' # training | validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "seg_files = [x for x in Path(f'./ADEChallengeData2016/annotations/{split}').glob('*.png')]\n",
    "seg_files.sort()\n",
    "\n",
    "seg_files = [[line_id+1, x] for line_id, x in enumerate(seg_files)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_row(line_id, seg_file):\n",
    "    stem = seg_file.stem\n",
    "    jpg_file = Path(f'./ADEChallengeData2016/images/{split}/{stem}.jpg')\n",
    "    \n",
    "    image = Image.open(jpg_file)\n",
    "\n",
    "    output = BytesIO()\n",
    "    image.save(output, 'PNG')\n",
    "    image_base64_str = base64.b64encode(output.getvalue())\n",
    "    image_base64_str = image_base64_str.decode(\"utf-8\")\n",
    "\n",
    "    seg = Image.open(seg_file)\n",
    "    seg = np.asarray(seg).copy()\n",
    "    mask_dict = {}\n",
    "    for seg_label in np.unique(seg.flatten()):\n",
    "        mask_dict[seg_label] = seg==seg_label\n",
    "\n",
    "    for seg_label, mask in mask_dict.items():\n",
    "        seg[mask] = shifted_id[seg_label]\n",
    "    \n",
    "    seg = Image.fromarray(seg)\n",
    "    output = BytesIO()\n",
    "    seg.save(output, 'PNG')\n",
    "    seg_base64_str = base64.b64encode(output.getvalue())\n",
    "    seg_base64_str = seg_base64_str.decode(\"utf-8\")\n",
    "    \n",
    "    img_id = stem[len('ADE_val_'):]\n",
    "    row = '\\t'.join([image_base64_str, seg_base64_str, img_id, str(line_id)])\n",
    "    \n",
    "    return row\n",
    "\n",
    "class Result():\n",
    "    def __init__(self):\n",
    "        self.rows = []\n",
    "\n",
    "    def update_result(self, row):\n",
    "        if row is not None:\n",
    "            self.rows.append(row)\n",
    "\n",
    "    def get_rows(self):\n",
    "        return self.rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool = Pool(128)\n",
    "result = Result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "row_list = []\n",
    "for line_id, seg_file in seg_files:\n",
    "    pool.apply_async(return_row, args=(line_id, seg_file), callback=result.update_result)\n",
    "pool.close()\n",
    "pool.join()\n",
    "\n",
    "rows = result.get_rows()\n",
    "rows.sort(key=lambda x: int(x.split('\\t')[-1]))\n",
    "ade_fullfile = \"\\n\".join(rows)\n",
    "\n",
    "with open(f'./dataset/ade/{split}.tsv', 'w') as f:\n",
    "    f.write(ade_fullfile)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.10 64-bit",
   "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.9.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}