{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "283f6e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import glob\n",
    "import json\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "path_dataset = r'D:\\Karusel\\Nurburg-karussel_88'\n",
    "with open(os.path.join(path_dataset, 'meta.json'), 'r') as j:\n",
    "    meta = json.load(j)\n",
    "\n",
    "imgs = glob.glob(path_dataset + '\\\\images\\\\img\\\\*', recursive=True)\n",
    "anns = glob.glob(path_dataset + '\\\\images\\\\ann\\\\*', recursive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7af948d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def label2hash(meta_json, last):\n",
    "    for clss in meta_json['classes']:\n",
    "        if clss['title'] == last['classTitle']:\n",
    "            meta_nodes = clss['geometry_config']['nodes']\n",
    "            label2hash = {}\n",
    "            for name in meta_nodes:\n",
    "                label2hash[meta_nodes[name]['label']] = name\n",
    "    return label2hash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7ea0771e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotations(meta_json, obj):\n",
    "    nodes = obj['nodes']\n",
    "    keypoints_2d = pd.DataFrame(columns=['x', 'y'])\n",
    "\n",
    "    for i in range(1, len(nodes)+1):\n",
    "        keypoints_2d.loc[i] = nodes[label2hash(meta_json, obj)[str(i)]]['loc']\n",
    "\n",
    "    keypoints_2d['v'] = 2\n",
    "    keypoints_2d[(keypoints_2d.x==0)&(keypoints_2d.y==0)] = 0\n",
    "    keypoints_2d = keypoints_2d.astype(float).round().astype(int)\n",
    "    return keypoints_2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "26b6abf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ann_json(keypoints, img_id, obj):\n",
    "    \n",
    "    annotation = {\n",
    "        \"id\": obj['id'],\n",
    "        \"segmentation\": [],\n",
    "        \"num_keypoints\": len(keypoints),\n",
    "        \"area\": 0,\n",
    "        \"iscrowd\": 0,\n",
    "        \"image_id\": img_id,\n",
    "        \"bbox\": [],\n",
    "        \"category_id\": 1,\n",
    "        \"keypoints\": keypoints.values.flatten().tolist()}\n",
    "\n",
    "    return annotation\n",
    "\n",
    "def img_json(ann, name, id):\n",
    "    height, width =  ann['size'].values()\n",
    "    image = {\n",
    "        \"id\": id,\n",
    "        \"width\": width,\n",
    "        \"height\": height,\n",
    "        \"file_name\": name,\n",
    "    }\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7c4d21a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def skeleton_supervisely(meta):\n",
    "    meta_nodes = meta['classes'][0]['geometry_config']['nodes']\n",
    "    label2hash = {}\n",
    "    for name in meta_nodes:\n",
    "        label2hash[meta_nodes[name]['label']] = name\n",
    "\n",
    "    skeleton_supervisely = []\n",
    "    for edge in meta['classes'][0]['geometry_config']['edges']:\n",
    "        skeleton_supervisely.append([meta_nodes[edge['src']]['label'], meta_nodes[edge['dst']]['label']])\n",
    "    return skeleton_supervisely"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "155be42d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ann_img_list(anns, imgs, meta):\n",
    "    annotations_list = []\n",
    "    image_list = []\n",
    "    for i in tqdm(range(len(anns))):\n",
    "\n",
    "        with open(anns[i], 'r') as j:\n",
    "            ann = json.load(j)\n",
    "        \n",
    "        image_name = os.path.basename(anns[i])[:-5]\n",
    "        image = img_json(ann, image_name, i)\n",
    "        image_list.append(image)\n",
    "\n",
    "        for obj in ann['objects']:\n",
    "            keypoints = annotations(meta, obj)\n",
    "            annotations_list.append(ann_json(keypoints, i, obj))\n",
    "    return image_list, annotations_list, meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6593f677",
   "metadata": {},
   "outputs": [],
   "source": [
    "def COCO(image_list, annotations_list, meta):\n",
    "    num_kpts = len(meta['classes'][0]['geometry_config']['nodes'])\n",
    "    coco = {\n",
    "\n",
    "        \"info\": {\n",
    "            \"description\": \"karusel Dataset\", \"version\": \"2.0\", \"keypoints\":num_kpts\n",
    "        },\n",
    "\n",
    "        \"categories\": [\n",
    "            {\n",
    "                \"supercategory\": \"NurburgRing\",\n",
    "                \"id\": 1,\n",
    "                \"name\": \"karusel\",\n",
    "                \"keypoints\": list(range(num_kpts)),\n",
    "                \"skeleton\": skeleton_supervisely(meta)\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "\n",
    "    coco['images'] = image_list\n",
    "    coco['annotations'] = annotations_list\n",
    "    return coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "38880d13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "219ba034d89346f188cf37dc9a722263",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/448 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "coco_json = COCO(*ann_img_list(anns, imgs, meta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "314565c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(path_dataset, 'karusel_COCO.json'), 'w') as file:\n",
    "    json.dump(coco_json, file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}