data_preprocessing.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# labelme2coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from glob import glob\n",
    "import mmcv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ann_path = './data/train/*.txt'\n",
    "ann_files = glob(ann_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def labelmetxt2coco(anns):\n",
    "    cars = {}\n",
    "    # 숫자로 해도 되지만, 클래스 확인하기 위해 차종으로 변경\n",
    "    classes = [\"chevrolet_malibu_sedan_2012_2016\", \"chevrolet_malibu_sedan_2017_2019\", \"chevrolet_spark_hatchback_2016_2021\", \"chevrolet_trailblazer_suv_2021_\", \"chevrolet_trax_suv_2017_2019\", \"genesis_g80_sedan_2016_2020\", \"genesis_g80_sedan_2021_\", \"genesis_gv80_suv_2020_\", \"hyundai_avante_sedan_2011_2015\", \"hyundai_avante_sedan_2020_\", \"hyundai_grandeur_sedan_2011_2016\", \"hyundai_grandstarex_van_2018_2020\", \"hyundai_ioniq_hatchback_2016_2019\", \"hyundai_sonata_sedan_2004_2009\", \"hyundai_sonata_sedan_2010_2014\", \"hyundai_sonata_sedan_2019_2020\", \"kia_carnival_van_2015_2020\", \"kia_carnival_van_2021_\", \"kia_k5_sedan_2010_2015\", \"kia_k5_sedan_2020_\", \"kia_k7_sedan_2016_2020\", \"kia_mohave_suv_2020_\", \"kia_morning_hatchback_2004_2010\", \"kia_morning_hatchback_2011_2016\", \"kia_ray_hatchback_2012_2017\", \"kia_sorrento_suv_2015_2019\", \"kia_sorrento_suv_2020_\", \"kia_soul_suv_2014_2018\", \"kia_sportage_suv_2016_2020\", \"kia_stonic_suv_2017_2019\", \"renault_sm3_sedan_2015_2018\", \"renault_xm3_suv_2020_\", \"ssangyong_korando_suv_2019_2020\", \"ssangyong_tivoli_suv_2016_2020\"]\n",
    "    cars[\"categories\"] = [{\"id\":i,\"name\":cat,\"supercategory\":\"none\"} for i, cat in enumerate(classes)]\n",
    "    cars[\"images\"] = []\n",
    "    cars[\"annotations\"] = []\n",
    "    cnt_ann = 0\n",
    "    for i, ann in enumerate(anns):\n",
    "        img_path = ann.replace(\"/data\", \"\").replace(\"txt\", \"png\")\n",
    "        #h, w, _ = cv2.imread(img_path).shape\n",
    "        cars[\"images\"].append({\"id\":i,\"height\":1040,\"width\":1920,\"file_name\":img_path})\n",
    "\n",
    "        f_ann = open(ann, \"r\")\n",
    "        for line in f_ann.readlines():\n",
    "            data = line.split()\n",
    "            cat = int(float(data[0]))\n",
    "            pt1x = int(data[1])\n",
    "            pt1y = int(data[2])\n",
    "            #pt2x = int(data[3])    not to use\n",
    "            #pt2y = int(data[4])\n",
    "            pt3x = int(data[5])\n",
    "            pt3y = int(data[6])\n",
    "            #pt4x = int(data[7])\n",
    "            #pt4y = int(data[8])\n",
    "            x = pt1x\n",
    "            y = pt1y\n",
    "            width = pt3x - pt1x\n",
    "            height = pt3y - pt1y\n",
    "            area = width * height\n",
    "            cars[\"annotations\"].append({\"id\": cnt_ann,\n",
    "                                        \"image_id\": i,\n",
    "                                        \"category_id\": cat,\n",
    "                                        \"bbox\": [x, y, width, height],\n",
    "                                        \"area\": area,\n",
    "                                        \"segmentation\": [],\n",
    "                                        \"iscrowd\": 0})\n",
    "            cnt_ann += 1\n",
    "        \n",
    "    return cars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/annotations/train.json', \"w\") as f:\n",
    "    json.dump(labelmetxt2coco(ann_files), f, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# test ann_file 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "test_img_path = './data/test/*'\n",
    "test_img_files = glob(test_img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset size: 3400\n"
     ]
    }
   ],
   "source": [
    "img_infos = []\n",
    "for idx, img_path in enumerate(test_img_files):\n",
    "    filename = img_path.replace(\"/data\", \"\")\n",
    "    height, width = mmcv.imread(str(img_path)).shape[:2]\n",
    "    img_infos.append(\n",
    "        dict(id=idx, file_name=filename, width=width, height=height))\n",
    "\n",
    "print(f'dataset size: {len(img_infos)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = {}\n",
    "classes = [\"chevrolet_malibu_sedan_2012_2016\", \"chevrolet_malibu_sedan_2017_2019\", \"chevrolet_spark_hatchback_2016_2021\", \"chevrolet_trailblazer_suv_2021_\", \"chevrolet_trax_suv_2017_2019\", \"genesis_g80_sedan_2016_2020\", \"genesis_g80_sedan_2021_\", \"genesis_gv80_suv_2020_\", \"hyundai_avante_sedan_2011_2015\", \"hyundai_avante_sedan_2020_\", \"hyundai_grandeur_sedan_2011_2016\", \"hyundai_grandstarex_van_2018_2020\", \"hyundai_ioniq_hatchback_2016_2019\", \"hyundai_sonata_sedan_2004_2009\", \"hyundai_sonata_sedan_2010_2014\", \"hyundai_sonata_sedan_2019_2020\", \"kia_carnival_van_2015_2020\", \"kia_carnival_van_2021_\", \"kia_k5_sedan_2010_2015\", \"kia_k5_sedan_2020_\", \"kia_k7_sedan_2016_2020\", \"kia_mohave_suv_2020_\", \"kia_morning_hatchback_2004_2010\", \"kia_morning_hatchback_2011_2016\", \"kia_ray_hatchback_2012_2017\", \"kia_sorrento_suv_2015_2019\", \"kia_sorrento_suv_2020_\", \"kia_soul_suv_2014_2018\", \"kia_sportage_suv_2016_2020\", \"kia_stonic_suv_2017_2019\", \"renault_sm3_sedan_2015_2018\", \"renault_xm3_suv_2020_\", \"ssangyong_korando_suv_2019_2020\", \"ssangyong_tivoli_suv_2016_2020\"]\n",
    "test[\"categories\"] = [{\"id\":i,\"name\":cat,\"supercategory\":\"none\"} for i, cat in enumerate(classes)]\n",
    "test[\"images\"] = img_infos\n",
    "\n",
    "with open('./data/annotations/test.json', \"w\") as f:\n",
    "    json.dump(test, f, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Split data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.model_selection import StratifiedGroupKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TRAIN: [   0    0    1 ... 6480 6480 6480]\n",
      "  [ 5 17 32 ... 25 17 21]\n",
      " TEST: [   2    2    2 ... 6479 6479 6479]\n",
      "  [23 22 24 ... 24 32 32]\n",
      "TRAIN: [   0    0    2 ... 6479 6479 6479]\n",
      "  [ 5 17 23 ... 24 32 32]\n",
      " TEST: [   1    1    1 ... 6480 6480 6480]\n",
      "  [32 26 24 ... 25 17 21]\n",
      "TRAIN: [   0    0    1 ... 6480 6480 6480]\n",
      "  [ 5 17 32 ... 25 17 21]\n",
      " TEST: [  10   10   10 ... 6473 6473 6473]\n",
      "  [16  3 23 ...  8 15 18]\n",
      "TRAIN: [   0    0    1 ... 6480 6480 6480]\n",
      "  [ 5 17 32 ... 25 17 21]\n",
      " TEST: [   7    7    7 ... 6478 6478 6478]\n",
      "  [ 3  1  6 ... 31 10 31]\n",
      "TRAIN: [   1    1    1 ... 6480 6480 6480]\n",
      "  [32 26 24 ... 25 17 21]\n",
      " TEST: [   0    0    4 ... 6475 6475 6475]\n",
      "  [ 5 17  5 ... 15 21 23]\n"
     ]
    }
   ],
   "source": [
    "# annotation = {train.json dataset file 경로}\n",
    "annotation = './data/annotations/train.json'\n",
    "\n",
    "with open(annotation) as f: data = json.load(f)\n",
    "\n",
    "var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]\n",
    "X = np.ones((len(data['annotations']),1))\n",
    "y = np.array([v[1] for v in var])\n",
    "groups = np.array([v[0] for v in var])\n",
    "\n",
    "cv = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=777)\n",
    "\n",
    "for train_idx, val_idx in cv.split(X, y, groups):\n",
    "    print(\"TRAIN:\", groups[train_idx]) # image_id\n",
    "    print(\" \", y[train_idx])           # category_id\n",
    "    print(\" TEST:\", groups[val_idx])\n",
    "    print(\" \", y[val_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_distribution(y):\n",
    "    y_distr = Counter(y)\n",
    "    y_vals_sum = sum(y_distr.values())\n",
    "\n",
    "    return [f'{y_distr[i]/y_vals_sum:.2%}' for i in range(np.max(y) +1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "distrs = [get_distribution(y)]\n",
    "index = ['training set']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(X, y, groups)):\n",
    "    train_y, val_y = y[train_idx], y[val_idx]\n",
    "    train_gr, val_gr = groups[train_idx], groups[val_idx]\n",
    "\n",
    "    assert len(set(train_gr) & set(val_gr)) == 0\n",
    "    \n",
    "    distrs.append(get_distribution(train_y))\n",
    "    distrs.append(get_distribution(val_y))\n",
    "    index.append(f'train - fold{fold_ind+1}')\n",
    "    index.append(f'val - fold{fold_ind+1}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>chevrolet_malibu_sedan_2012_2016</th>\n",
       "      <th>chevrolet_malibu_sedan_2017_2019</th>\n",
       "      <th>chevrolet_spark_hatchback_2016_2021</th>\n",
       "      <th>chevrolet_trailblazer_suv_2021_</th>\n",
       "      <th>chevrolet_trax_suv_2017_2019</th>\n",
       "      <th>genesis_g80_sedan_2016_2020</th>\n",
       "      <th>genesis_g80_sedan_2021_</th>\n",
       "      <th>genesis_gv80_suv_2020_</th>\n",
       "      <th>hyundai_avante_sedan_2011_2015</th>\n",
       "      <th>hyundai_avante_sedan_2020_</th>\n",
       "      <th>...</th>\n",
       "      <th>kia_ray_hatchback_2012_2017</th>\n",
       "      <th>kia_sorrento_suv_2015_2019</th>\n",
       "      <th>kia_sorrento_suv_2020_</th>\n",
       "      <th>kia_soul_suv_2014_2018</th>\n",
       "      <th>kia_sportage_suv_2016_2020</th>\n",
       "      <th>kia_stonic_suv_2017_2019</th>\n",
       "      <th>renault_sm3_sedan_2015_2018</th>\n",
       "      <th>renault_xm3_suv_2020_</th>\n",
       "      <th>ssangyong_korando_suv_2019_2020</th>\n",
       "      <th>ssangyong_tivoli_suv_2016_2020</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>training set</th>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold1</th>\n",
       "      <td>3.01%</td>\n",
       "      <td>3.04%</td>\n",
       "      <td>3.04%</td>\n",
       "      <td>2.84%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.87%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.80%</td>\n",
       "      <td>2.84%</td>\n",
       "      <td>2.78%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.88%</td>\n",
       "      <td>2.96%</td>\n",
       "      <td>2.82%</td>\n",
       "      <td>2.96%</td>\n",
       "      <td>2.94%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold1</th>\n",
       "      <td>2.67%</td>\n",
       "      <td>2.56%</td>\n",
       "      <td>2.56%</td>\n",
       "      <td>3.35%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.58%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>3.23%</td>\n",
       "      <td>...</td>\n",
       "      <td>3.49%</td>\n",
       "      <td>3.35%</td>\n",
       "      <td>3.58%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.79%</td>\n",
       "      <td>3.17%</td>\n",
       "      <td>2.85%</td>\n",
       "      <td>3.44%</td>\n",
       "      <td>2.88%</td>\n",
       "      <td>2.94%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold2</th>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.76%</td>\n",
       "      <td>3.00%</td>\n",
       "      <td>2.96%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>2.87%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.93%</td>\n",
       "      <td>2.93%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>3.06%</td>\n",
       "      <td>3.00%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.88%</td>\n",
       "      <td>2.94%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold2</th>\n",
       "      <td>2.80%</td>\n",
       "      <td>3.69%</td>\n",
       "      <td>2.71%</td>\n",
       "      <td>2.86%</td>\n",
       "      <td>2.77%</td>\n",
       "      <td>3.07%</td>\n",
       "      <td>3.24%</td>\n",
       "      <td>2.59%</td>\n",
       "      <td>3.04%</td>\n",
       "      <td>3.01%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.62%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.59%</td>\n",
       "      <td>2.45%</td>\n",
       "      <td>2.71%</td>\n",
       "      <td>3.04%</td>\n",
       "      <td>3.18%</td>\n",
       "      <td>2.95%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold3</th>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.86%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>2.85%</td>\n",
       "      <td>3.04%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>...</td>\n",
       "      <td>3.05%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>3.06%</td>\n",
       "      <td>2.96%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.93%</td>\n",
       "      <td>3.00%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.99%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold3</th>\n",
       "      <td>2.94%</td>\n",
       "      <td>3.26%</td>\n",
       "      <td>3.09%</td>\n",
       "      <td>2.62%</td>\n",
       "      <td>2.62%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>3.12%</td>\n",
       "      <td>3.29%</td>\n",
       "      <td>2.53%</td>\n",
       "      <td>3.12%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.50%</td>\n",
       "      <td>3.06%</td>\n",
       "      <td>2.47%</td>\n",
       "      <td>2.88%</td>\n",
       "      <td>2.94%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>3.00%</td>\n",
       "      <td>2.71%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>2.74%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold4</th>\n",
       "      <td>2.84%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.81%</td>\n",
       "      <td>2.99%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>2.97%</td>\n",
       "      <td>2.99%</td>\n",
       "      <td>2.89%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.97%</td>\n",
       "      <td>3.05%</td>\n",
       "      <td>3.01%</td>\n",
       "      <td>2.86%</td>\n",
       "      <td>2.86%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>3.00%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>2.92%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold4</th>\n",
       "      <td>3.34%</td>\n",
       "      <td>2.61%</td>\n",
       "      <td>3.46%</td>\n",
       "      <td>2.75%</td>\n",
       "      <td>3.08%</td>\n",
       "      <td>2.81%</td>\n",
       "      <td>2.75%</td>\n",
       "      <td>3.14%</td>\n",
       "      <td>3.11%</td>\n",
       "      <td>2.78%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.84%</td>\n",
       "      <td>2.49%</td>\n",
       "      <td>2.67%</td>\n",
       "      <td>3.25%</td>\n",
       "      <td>3.28%</td>\n",
       "      <td>3.02%</td>\n",
       "      <td>3.08%</td>\n",
       "      <td>2.72%</td>\n",
       "      <td>2.64%</td>\n",
       "      <td>3.02%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold5</th>\n",
       "      <td>2.94%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.89%</td>\n",
       "      <td>2.87%</td>\n",
       "      <td>2.93%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>2.99%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.88%</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>3.03%</td>\n",
       "      <td>2.90%</td>\n",
       "      <td>2.89%</td>\n",
       "      <td>2.91%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>2.93%</td>\n",
       "      <td>2.91%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold5</th>\n",
       "      <td>2.95%</td>\n",
       "      <td>2.60%</td>\n",
       "      <td>2.89%</td>\n",
       "      <td>3.13%</td>\n",
       "      <td>3.22%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>3.01%</td>\n",
       "      <td>2.74%</td>\n",
       "      <td>3.10%</td>\n",
       "      <td>2.57%</td>\n",
       "      <td>...</td>\n",
       "      <td>2.92%</td>\n",
       "      <td>3.19%</td>\n",
       "      <td>3.01%</td>\n",
       "      <td>2.57%</td>\n",
       "      <td>3.10%</td>\n",
       "      <td>3.16%</td>\n",
       "      <td>3.07%</td>\n",
       "      <td>2.80%</td>\n",
       "      <td>2.98%</td>\n",
       "      <td>3.07%</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>11 rows × 34 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              chevrolet_malibu_sedan_2012_2016  \\\n",
       "training set                             2.94%   \n",
       "train - fold1                            3.01%   \n",
       "val - fold1                              2.67%   \n",
       "train - fold2                            2.98%   \n",
       "val - fold2                              2.80%   \n",
       "train - fold3                            2.94%   \n",
       "val - fold3                              2.94%   \n",
       "train - fold4                            2.84%   \n",
       "val - fold4                              3.34%   \n",
       "train - fold5                            2.94%   \n",
       "val - fold5                              2.95%   \n",
       "\n",
       "              chevrolet_malibu_sedan_2017_2019  \\\n",
       "training set                             2.94%   \n",
       "train - fold1                            3.04%   \n",
       "val - fold1                              2.56%   \n",
       "train - fold2                            2.76%   \n",
       "val - fold2                              3.69%   \n",
       "train - fold3                            2.86%   \n",
       "val - fold3                              3.26%   \n",
       "train - fold4                            3.02%   \n",
       "val - fold4                              2.61%   \n",
       "train - fold5                            3.03%   \n",
       "val - fold5                              2.60%   \n",
       "\n",
       "              chevrolet_spark_hatchback_2016_2021  \\\n",
       "training set                                2.94%   \n",
       "train - fold1                               3.04%   \n",
       "val - fold1                                 2.56%   \n",
       "train - fold2                               3.00%   \n",
       "val - fold2                                 2.71%   \n",
       "train - fold3                               2.90%   \n",
       "val - fold3                                 3.09%   \n",
       "train - fold4                               2.81%   \n",
       "val - fold4                                 3.46%   \n",
       "train - fold5                               2.95%   \n",
       "val - fold5                                 2.89%   \n",
       "\n",
       "              chevrolet_trailblazer_suv_2021_ chevrolet_trax_suv_2017_2019  \\\n",
       "training set                            2.94%                        2.94%   \n",
       "train - fold1                           2.84%                        2.92%   \n",
       "val - fold1                             3.35%                        3.02%   \n",
       "train - fold2                           2.96%                        2.98%   \n",
       "val - fold2                             2.86%                        2.77%   \n",
       "train - fold3                           3.02%                        3.02%   \n",
       "val - fold3                             2.62%                        2.62%   \n",
       "train - fold4                           2.99%                        2.91%   \n",
       "val - fold4                             2.75%                        3.08%   \n",
       "train - fold5                           2.89%                        2.87%   \n",
       "val - fold5                             3.13%                        3.22%   \n",
       "\n",
       "              genesis_g80_sedan_2016_2020 genesis_g80_sedan_2021_  \\\n",
       "training set                        2.94%                   2.94%   \n",
       "train - fold1                       2.94%                   3.03%   \n",
       "val - fold1                         2.94%                   2.58%   \n",
       "train - fold2                       2.91%                   2.87%   \n",
       "val - fold2                         3.07%                   3.24%   \n",
       "train - fold3                       2.95%                   2.90%   \n",
       "val - fold3                         2.91%                   3.12%   \n",
       "train - fold4                       2.97%                   2.99%   \n",
       "val - fold4                         2.81%                   2.75%   \n",
       "train - fold5                       2.93%                   2.92%   \n",
       "val - fold5                         2.98%                   3.01%   \n",
       "\n",
       "              genesis_gv80_suv_2020_ hyundai_avante_sedan_2011_2015  \\\n",
       "training set                   2.94%                          2.94%   \n",
       "train - fold1                  2.94%                          2.94%   \n",
       "val - fold1                    2.94%                          2.94%   \n",
       "train - fold2                  3.03%                          2.92%   \n",
       "val - fold2                    2.59%                          3.04%   \n",
       "train - fold3                  2.85%                          3.04%   \n",
       "val - fold3                    3.29%                          2.53%   \n",
       "train - fold4                  2.89%                          2.90%   \n",
       "val - fold4                    3.14%                          3.11%   \n",
       "train - fold5                  2.99%                          2.90%   \n",
       "val - fold5                    2.74%                          3.10%   \n",
       "\n",
       "              hyundai_avante_sedan_2020_  ... kia_ray_hatchback_2012_2017  \\\n",
       "training set                       2.94%  ...                       2.94%   \n",
       "train - fold1                      2.87%  ...                       2.80%   \n",
       "val - fold1                        3.23%  ...                       3.49%   \n",
       "train - fold2                      2.92%  ...                       2.94%   \n",
       "val - fold2                        3.01%  ...                       2.95%   \n",
       "train - fold3                      2.90%  ...                       3.05%   \n",
       "val - fold3                        3.12%  ...                       2.50%   \n",
       "train - fold4                      2.98%  ...                       2.97%   \n",
       "val - fold4                        2.78%  ...                       2.84%   \n",
       "train - fold5                      3.03%  ...                       2.95%   \n",
       "val - fold5                        2.57%  ...                       2.92%   \n",
       "\n",
       "              kia_sorrento_suv_2015_2019 kia_sorrento_suv_2020_  \\\n",
       "training set                       2.94%                  2.94%   \n",
       "train - fold1                      2.84%                  2.78%   \n",
       "val - fold1                        3.35%                  3.58%   \n",
       "train - fold2                      3.02%                  2.93%   \n",
       "val - fold2                        2.62%                  2.98%   \n",
       "train - fold3                      2.91%                  3.06%   \n",
       "val - fold3                        3.06%                  2.47%   \n",
       "train - fold4                      3.05%                  3.01%   \n",
       "val - fold4                        2.49%                  2.67%   \n",
       "train - fold5                      2.88%                  2.92%   \n",
       "val - fold5                        3.19%                  3.01%   \n",
       "\n",
       "              kia_soul_suv_2014_2018 kia_sportage_suv_2016_2020  \\\n",
       "training set                   2.94%                      2.94%   \n",
       "train - fold1                  2.92%                      2.98%   \n",
       "val - fold1                    3.02%                      2.79%   \n",
       "train - fold2                  2.93%                      3.03%   \n",
       "val - fold2                    2.98%                      2.59%   \n",
       "train - fold3                  2.96%                      2.94%   \n",
       "val - fold3                    2.88%                      2.94%   \n",
       "train - fold4                  2.86%                      2.86%   \n",
       "val - fold4                    3.25%                      3.28%   \n",
       "train - fold5                  3.03%                      2.90%   \n",
       "val - fold5                    2.57%                      3.10%   \n",
       "\n",
       "              kia_stonic_suv_2017_2019 renault_sm3_sedan_2015_2018  \\\n",
       "training set                     2.94%                       2.94%   \n",
       "train - fold1                    2.88%                       2.96%   \n",
       "val - fold1                      3.17%                       2.85%   \n",
       "train - fold2                    3.06%                       3.00%   \n",
       "val - fold2                      2.45%                       2.71%   \n",
       "train - fold3                    2.95%                       2.93%   \n",
       "val - fold3                      2.91%                       3.00%   \n",
       "train - fold4                    2.92%                       2.91%   \n",
       "val - fold4                      3.02%                       3.08%   \n",
       "train - fold5                    2.89%                       2.91%   \n",
       "val - fold5                      3.16%                       3.07%   \n",
       "\n",
       "              renault_xm3_suv_2020_ ssangyong_korando_suv_2019_2020  \\\n",
       "training set                  2.94%                           2.94%   \n",
       "train - fold1                 2.82%                           2.96%   \n",
       "val - fold1                   3.44%                           2.88%   \n",
       "train - fold2                 2.92%                           2.88%   \n",
       "val - fold2                   3.04%                           3.18%   \n",
       "train - fold3                 3.00%                           2.92%   \n",
       "val - fold3                   2.71%                           3.03%   \n",
       "train - fold4                 3.00%                           3.02%   \n",
       "val - fold4                   2.72%                           2.64%   \n",
       "train - fold5                 2.98%                           2.93%   \n",
       "val - fold5                   2.80%                           2.98%   \n",
       "\n",
       "              ssangyong_tivoli_suv_2016_2020  \n",
       "training set                           2.94%  \n",
       "train - fold1                          2.94%  \n",
       "val - fold1                            2.94%  \n",
       "train - fold2                          2.94%  \n",
       "val - fold2                            2.95%  \n",
       "train - fold3                          2.99%  \n",
       "val - fold3                            2.74%  \n",
       "train - fold4                          2.92%  \n",
       "val - fold4                            3.02%  \n",
       "train - fold5                          2.91%  \n",
       "val - fold5                            3.07%  \n",
       "\n",
       "[11 rows x 34 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "categories = [d['name'] for d in data['categories']]\n",
    "pd.DataFrame(distrs, index=index, columns = [categories[i] for i in range(np.max(y) + 1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# annotation = {dataset 경로/K-fold}\n",
    "output_filename = \"./data/K-fold\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, (train_idx, val_idx) in enumerate(cv.split(X, y, groups)):\n",
    "    train_images, val_images = [], []\n",
    "    train_annotations, val_annotations = [], []\n",
    "    for i in groups[train_idx]: # image_id\n",
    "        train_images.append(data[\"images\"][i].copy())\n",
    "    for i in groups[val_idx]:   # image_id\n",
    "        val_images.append(data[\"images\"][i].copy())\n",
    "    for annotation in data[\"annotations\"]:\n",
    "        if annotation[\"image_id\"] in groups[val_idx]:\n",
    "            val_annotations.append(annotation.copy())\n",
    "        else:\n",
    "            train_annotations.append(annotation.copy())\n",
    "\n",
    "    train_split = {\n",
    "            \"images\": train_images,\n",
    "            \"annotations\": train_annotations,\n",
    "            \"info\": data.get(\"info\", {}),\n",
    "            \"licenses\": data.get(\"licenses\", []),\n",
    "            \"categories\": data[\"categories\"],\n",
    "        }\n",
    "\n",
    "    val_split = {\n",
    "            \"images\": val_images,\n",
    "            \"annotations\": val_annotations,\n",
    "            \"info\": data.get(\"info\", {}),\n",
    "            \"licenses\": data.get(\"licenses\", []),\n",
    "            \"categories\": data[\"categories\"],\n",
    "        }\n",
    "    \n",
    "    output_files = []\n",
    "    for split_type, split in zip([\"train\", \"val\"], [train_split, val_split]):\n",
    "        output_files.append(output_filename + f\"_{split_type}{idx+1}.json\")\n",
    "        with open(output_files[-1], \"w\") as f:\n",
    "            json.dump(split, f, indent=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dacon",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}