convert_segmentation_coco.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_id = {\n",
    "    0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 12: 11, 13: 12, 14: 13, 15: 14, 16: 15, 17: 16, 18: 17, 19: 18, 20: 19, 21: 20, 22: 21,\n",
    "    23: 22, 24: 23, 26: 24, 27: 25, 30: 26, 31: 27, 32: 28, 33: 29, 34: 30, 35: 31, 36: 32, 37: 33, 38: 34, 39: 35, 40: 36, 41: 37, 42: 38, 43: 39, 45: 40, 46: 41,\n",
    "    47: 42, 48: 43, 49: 44, 50: 45, 51: 46, 52: 47, 53: 48, 54: 49, 55: 50, 56: 51, 57: 52, 58: 53, 59: 54, 60: 55, 61: 56, 62: 57, 63: 58, 64: 59, 66: 60, 69: 61,\n",
    "    71: 62, 72: 63, 73: 64, 74: 65, 75: 66, 76: 67, 77: 68, 78: 69, 79: 70, 80: 71, 81: 72, 83: 73, 84: 74, 85: 75, 86: 76, 87: 77, 88: 78, 89: 79, 91: 80, 92: 81,\n",
    "    93: 82, 94: 83, 95: 84, 96: 85, 97: 86, 98: 87, 99: 88, 100: 89, 101: 90, 102: 91, 103: 92, 104: 93, 105: 94, 106: 95, 107: 96, 108: 97, 109: 98, 110: 99, 111: 100,\n",
    "    112: 101, 113: 102, 114: 103, 115: 104, 116: 105, 117: 106, 118: 107, 119: 108, 120: 109, 121: 110, 122: 111, 123: 112, 124: 113, 125: 114, 126: 115, 127: 116, 128: 117,\n",
    "    129: 118, 130: 119, 131: 120, 132: 121, 133: 122, 134: 123, 135: 124, 136: 125, 137: 126, 138: 127, 139: 128, 140: 129, 141: 130, 142: 131, 143: 132, 144: 133, 145: 134,\n",
    "    146: 135, 147: 136, 148: 137, 149: 138, 150: 139, 151: 140, 152: 141, 153: 142, 154: 143, 155: 144, 156: 145, 157: 146, 158: 147, 159: 148, 160: 149, 161: 150, 162: 151,\n",
    "    163: 152, 164: 153, 165: 154, 166: 155, 167: 156, 168: 157, 169: 158, 170: 159, 171: 160, 172: 161, 173: 162, 174: 163, 175: 164, 176: 165, 177: 166, 178: 167, 179: 168,\n",
    "    180: 169, 181: 170, 255: 0\n",
    "}\n",
    "\n",
    "shifted_fine_id = {}\n",
    "for key, val in fine_id.items():\n",
    "    val = val + 1\n",
    "    shifted_fine_id[key] = val\n",
    "shifted_fine_id[255] = 0\n",
    "fine_id = shifted_fine_id\n",
    "\n",
    "for key in range(182):\n",
    "    if key not in fine_id:\n",
    "        fine_id[key] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = 'val2017' # train2017 | val2017"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "fineseg_files = [x for x in Path(f'./COCO/labels/{split}').glob('*.png')]\n",
    "fineseg_files.sort()\n",
    "\n",
    "fineseg_files = [[line_id+1, x] for line_id, x in enumerate(fineseg_files)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_row(line_id, fineseg_file):\n",
    "    stem = fineseg_file.stem\n",
    "    jpg_file = Path(f'./COCO/{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",
    "    fineseg = Image.open(fineseg_file)\n",
    "    fineseg_refined = np.asarray(fineseg).copy()\n",
    "    mask_dict = {}\n",
    "    for seg_label in np.unique(fineseg_refined.flatten()):\n",
    "        mask_dict[seg_label] = fineseg_refined==seg_label\n",
    "\n",
    "    for seg_label, mask in mask_dict.items():\n",
    "        fineseg_refined[mask] = fine_id[seg_label]\n",
    "        \n",
    "    fineseg_refined = Image.fromarray(fineseg_refined)\n",
    "    output = BytesIO()\n",
    "    fineseg_refined.save(output, 'PNG')\n",
    "    fineseg_refined_base64_str = base64.b64encode(output.getvalue())\n",
    "    fineseg_refined_base64_str = fineseg_refined_base64_str.decode(\"utf-8\")\n",
    "\n",
    "    fineseg_refined_row = '\\t'.join([image_base64_str, fineseg_refined_base64_str, stem, str(line_id)])\n",
    "    \n",
    "    return fineseg_refined_row\n",
    "\n",
    "class Result():\n",
    "    def __init__(self):\n",
    "        self.fineseg_refined_rows = []\n",
    "\n",
    "    def update_result(self, fineseg_refined_row):\n",
    "        if fineseg_refined_row is not None:\n",
    "            self.fineseg_refined_rows.append(fineseg_refined_row)\n",
    "\n",
    "    def get_rows(self):\n",
    "        return self.fineseg_refined_rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool = Pool(128)\n",
    "result = Result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "row_list = []\n",
    "for line_id, fineseg_file in fineseg_files:\n",
    "    pool.apply_async(return_row, args=(line_id, fineseg_file), callback=result.update_result)\n",
    "pool.close()\n",
    "pool.join()\n",
    "\n",
    "fineseg_refined_rows = result.get_rows()\n",
    "fineseg_refined_rows.sort(key=lambda x: int(x.split('\\t')[-1]))\n",
    "fineseg_refined_rows_fullfile = \"\\n\".join(fineseg_refined_rows)\n",
    "\n",
    "with open(f'./dataset/coco/fineseg_refined_{split}.tsv', 'w') as f:\n",
    "    f.write(fineseg_refined_rows_fullfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}