{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f5321831-268f-43e2-8719-be399be625b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/botsort/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "import os\n",
    "import time\n",
    "import natsort\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9733b66a-88d8-45b5-9aee-c238d2b129f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = '/data/cvprw/AIC23/dataset/coco_AIC23/train'\n",
    "\n",
    "site_list = [file for file in os.listdir(dataset_path)]\n",
    "\n",
    "temp = []\n",
    "for site in site_list:\n",
    "    d = os.path.join(dataset_path, site)\n",
    "    if os.path.isdir(d) and not site.startswith('.'):\n",
    "#         print(channel)\n",
    "        temp.append(site)\n",
    "\n",
    "site_list = natsort.natsorted(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ea0330b3-ed44-4012-8f8e-01ab248eaeab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['S002',\n",
       " 'S004',\n",
       " 'S006',\n",
       " 'S007',\n",
       " 'S010',\n",
       " 'S011',\n",
       " 'S012',\n",
       " 'S015',\n",
       " 'S016',\n",
       " 'S019']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "site_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bfcedc65-af9e-4724-a76f-0ba719a1beb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# categoires part of coco\n",
    "category_list = ['person','none']\n",
    "category_dic_list=[]\n",
    "\n",
    "for idx, label in enumerate(category_list):\n",
    "    category_dic={}\n",
    "    category_dic[\"id\"]=idx+1\n",
    "    category_dic[\"name\"]=label\n",
    "    category_dic[\"supercategory\"]='AIC23'\n",
    "    category_dic_list.append(category_dic)\n",
    "    \n",
    "    \n",
    "    \n",
    "coco_format={}\n",
    "coco_format['images']=[]\n",
    "coco_format['annotations']=[]\n",
    "coco_format['categories']=category_dic_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "245ad45c-c11f-4788-885f-ed228d3d8b6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S002\n",
      "S004\n",
      "S006\n",
      "S007\n",
      "S010\n",
      "S011\n",
      "S012\n",
      "S015\n",
      "S016\n",
      "S019\n"
     ]
    }
   ],
   "source": [
    "img_id_count = 0\n",
    "ann_id_count = 0\n",
    "\n",
    "for site_1 in site_list:\n",
    "    print(site_1)\n",
    "    \n",
    "    file_path = os.path.join(\"/data/cvprw/AIC23/dataset/coco_AIC23/train\",site_1,\"train.json\")\n",
    "    \n",
    "    # returns JSON object as a dictionary\n",
    "    json_data = json.load(open(file_path,'r'))\n",
    "    \n",
    "    # print(len(json_data['images']))\n",
    "    # print(len(json_data['annotations']))\n",
    "    \n",
    "    for idx,img in enumerate(json_data['images']):\n",
    "        json_data['images'][idx]['id']+=img_id_count\n",
    "        coco_format['images'].append(json_data['images'][idx])\n",
    "\n",
    "    for idx,ann in enumerate(json_data['annotations']):\n",
    "        json_data['annotations'][idx]['image_id']+=img_id_count\n",
    "        json_data['annotations'][idx]['id']+=ann_id_count\n",
    "        coco_format['annotations'].append(json_data['annotations'][idx])\n",
    "        \n",
    "    img_id_count += len(json_data['images'])\n",
    "    ann_id_count += len(json_data['annotations'])\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6ecc116e-9ff9-4fa1-adc1-ef2ba6a325c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_path = dataset_path\n",
    "\n",
    "#save json\n",
    "with open(os.path.join(save_path,'train.json'), 'w', encoding='utf-8') as make_file:\n",
    "    json.dump(coco_format, make_file, indent=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e50a3fe-301c-4dd9-986f-e8655f944faf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "botsort",
   "language": "python",
   "name": "botsort"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}