create_valset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b9bd4ded-758d-4538-98b3-89b06fa1cf29",
   "metadata": {},
   "source": [
    "### StratifiedGroupKFold를 이용해 validation set 생성\n",
    "\n",
    "- 대회가 어느정도 막바지를 향해가면, n_split을  큰 값으로 늘려, OOF를 진행해볼 수 있을 것 같다"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "27bbfabd-f7a8-4db8-9fe6-ea053d0cfdb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TRAIN: [   0    1    1 ... 4881 4881 4881]\n",
      "  [0 3 7 ... 7 1 7]\n",
      " TEST: [   6   13   13 ... 4882 4882 4882]\n",
      "  [1 6 7 ... 0 1 1]\n",
      "TRAIN: [   0    1    1 ... 4882 4882 4882]\n",
      "  [0 3 7 ... 0 1 1]\n",
      " TEST: [   5    5    5 ... 4876 4876 4878]\n",
      "  [7 0 0 ... 0 2 0]\n",
      "TRAIN: [   0    3    3 ... 4882 4882 4882]\n",
      "  [0 2 6 ... 0 1 1]\n",
      " TEST: [   1    1    1 ... 4877 4877 4880]\n",
      "  [3 7 4 ... 7 7 0]\n",
      "TRAIN: [   1    1    1 ... 4882 4882 4882]\n",
      "  [3 7 4 ... 0 1 1]\n",
      " TEST: [   0    3    3 ... 4881 4881 4881]\n",
      "  [0 2 6 ... 7 1 7]\n",
      "TRAIN: [   0    1    1 ... 4882 4882 4882]\n",
      "  [0 3 7 ... 0 1 1]\n",
      " TEST: [   4    4    4 ... 4868 4872 4872]\n",
      "  [1 1 1 ... 2 4 6]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from sklearn.model_selection import StratifiedGroupKFold\n",
    "\n",
    "# load json \n",
    "annotation = 'dataset/train.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]\n",
    "\n",
    "X = np.ones((len(data['annotations']),1))\n",
    "y = np.array([v[1] for v in var])  # class\n",
    "groups = np.array([v[0] for v in var])  # image(group)\n",
    "\n",
    "cv = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=411)\n",
    "\n",
    "for train_idx, val_idx in cv.split(X, y, groups):\n",
    "    print(\"TRAIN:\", groups[train_idx])\n",
    "    print(\" \", y[train_idx])\n",
    "    print(\" TEST:\", groups[val_idx])\n",
    "    print(\" \", y[val_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "bfe3427f-45cf-4cb7-8ed4-766ea02f44e8",
   "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>General trash</th>\n",
       "      <th>Paper</th>\n",
       "      <th>Paper pack</th>\n",
       "      <th>Metal</th>\n",
       "      <th>Glass</th>\n",
       "      <th>Plastic</th>\n",
       "      <th>Styrofoam</th>\n",
       "      <th>Plastic bag</th>\n",
       "      <th>Battery</th>\n",
       "      <th>Clothing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>training set</th>\n",
       "      <td>17.14%</td>\n",
       "      <td>27.45%</td>\n",
       "      <td>3.88%</td>\n",
       "      <td>4.04%</td>\n",
       "      <td>4.24%</td>\n",
       "      <td>12.72%</td>\n",
       "      <td>5.46%</td>\n",
       "      <td>22.37%</td>\n",
       "      <td>0.69%</td>\n",
       "      <td>2.02%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold0</th>\n",
       "      <td>16.96%</td>\n",
       "      <td>27.45%</td>\n",
       "      <td>3.79%</td>\n",
       "      <td>4.13%</td>\n",
       "      <td>4.48%</td>\n",
       "      <td>12.61%</td>\n",
       "      <td>5.51%</td>\n",
       "      <td>22.28%</td>\n",
       "      <td>0.77%</td>\n",
       "      <td>2.02%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold0</th>\n",
       "      <td>17.85%</td>\n",
       "      <td>27.42%</td>\n",
       "      <td>4.23%</td>\n",
       "      <td>3.70%</td>\n",
       "      <td>3.26%</td>\n",
       "      <td>13.15%</td>\n",
       "      <td>5.25%</td>\n",
       "      <td>22.77%</td>\n",
       "      <td>0.35%</td>\n",
       "      <td>2.02%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold1</th>\n",
       "      <td>17.14%</td>\n",
       "      <td>27.24%</td>\n",
       "      <td>4.01%</td>\n",
       "      <td>3.98%</td>\n",
       "      <td>4.28%</td>\n",
       "      <td>12.77%</td>\n",
       "      <td>5.38%</td>\n",
       "      <td>22.32%</td>\n",
       "      <td>0.67%</td>\n",
       "      <td>2.20%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold1</th>\n",
       "      <td>17.12%</td>\n",
       "      <td>28.17%</td>\n",
       "      <td>3.41%</td>\n",
       "      <td>4.26%</td>\n",
       "      <td>4.12%</td>\n",
       "      <td>12.51%</td>\n",
       "      <td>5.72%</td>\n",
       "      <td>22.57%</td>\n",
       "      <td>0.73%</td>\n",
       "      <td>1.38%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold2</th>\n",
       "      <td>17.31%</td>\n",
       "      <td>27.39%</td>\n",
       "      <td>3.83%</td>\n",
       "      <td>4.08%</td>\n",
       "      <td>4.13%</td>\n",
       "      <td>12.80%</td>\n",
       "      <td>5.14%</td>\n",
       "      <td>22.68%</td>\n",
       "      <td>0.69%</td>\n",
       "      <td>1.94%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold2</th>\n",
       "      <td>16.42%</td>\n",
       "      <td>27.68%</td>\n",
       "      <td>4.05%</td>\n",
       "      <td>3.88%</td>\n",
       "      <td>4.70%</td>\n",
       "      <td>12.36%</td>\n",
       "      <td>6.76%</td>\n",
       "      <td>21.12%</td>\n",
       "      <td>0.69%</td>\n",
       "      <td>2.35%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold3</th>\n",
       "      <td>17.30%</td>\n",
       "      <td>27.47%</td>\n",
       "      <td>3.87%</td>\n",
       "      <td>4.06%</td>\n",
       "      <td>4.22%</td>\n",
       "      <td>12.63%</td>\n",
       "      <td>5.49%</td>\n",
       "      <td>22.39%</td>\n",
       "      <td>0.63%</td>\n",
       "      <td>1.95%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold3</th>\n",
       "      <td>16.50%</td>\n",
       "      <td>27.36%</td>\n",
       "      <td>3.88%</td>\n",
       "      <td>3.99%</td>\n",
       "      <td>4.33%</td>\n",
       "      <td>13.07%</td>\n",
       "      <td>5.33%</td>\n",
       "      <td>22.30%</td>\n",
       "      <td>0.92%</td>\n",
       "      <td>2.32%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold4</th>\n",
       "      <td>16.97%</td>\n",
       "      <td>27.67%</td>\n",
       "      <td>3.88%</td>\n",
       "      <td>3.97%</td>\n",
       "      <td>4.10%</td>\n",
       "      <td>12.77%</td>\n",
       "      <td>5.76%</td>\n",
       "      <td>22.20%</td>\n",
       "      <td>0.68%</td>\n",
       "      <td>2.00%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold4</th>\n",
       "      <td>17.84%</td>\n",
       "      <td>26.48%</td>\n",
       "      <td>3.85%</td>\n",
       "      <td>4.38%</td>\n",
       "      <td>4.84%</td>\n",
       "      <td>12.50%</td>\n",
       "      <td>4.15%</td>\n",
       "      <td>23.11%</td>\n",
       "      <td>0.73%</td>\n",
       "      <td>2.11%</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              General trash   Paper Paper pack  Metal  Glass Plastic  \\\n",
       "training set         17.14%  27.45%      3.88%  4.04%  4.24%  12.72%   \n",
       "train - fold0        16.96%  27.45%      3.79%  4.13%  4.48%  12.61%   \n",
       "val - fold0          17.85%  27.42%      4.23%  3.70%  3.26%  13.15%   \n",
       "train - fold1        17.14%  27.24%      4.01%  3.98%  4.28%  12.77%   \n",
       "val - fold1          17.12%  28.17%      3.41%  4.26%  4.12%  12.51%   \n",
       "train - fold2        17.31%  27.39%      3.83%  4.08%  4.13%  12.80%   \n",
       "val - fold2          16.42%  27.68%      4.05%  3.88%  4.70%  12.36%   \n",
       "train - fold3        17.30%  27.47%      3.87%  4.06%  4.22%  12.63%   \n",
       "val - fold3          16.50%  27.36%      3.88%  3.99%  4.33%  13.07%   \n",
       "train - fold4        16.97%  27.67%      3.88%  3.97%  4.10%  12.77%   \n",
       "val - fold4          17.84%  26.48%      3.85%  4.38%  4.84%  12.50%   \n",
       "\n",
       "              Styrofoam Plastic bag Battery Clothing  \n",
       "training set      5.46%      22.37%   0.69%    2.02%  \n",
       "train - fold0     5.51%      22.28%   0.77%    2.02%  \n",
       "val - fold0       5.25%      22.77%   0.35%    2.02%  \n",
       "train - fold1     5.38%      22.32%   0.67%    2.20%  \n",
       "val - fold1       5.72%      22.57%   0.73%    1.38%  \n",
       "train - fold2     5.14%      22.68%   0.69%    1.94%  \n",
       "val - fold2       6.76%      21.12%   0.69%    2.35%  \n",
       "train - fold3     5.49%      22.39%   0.63%    1.95%  \n",
       "val - fold3       5.33%      22.30%   0.92%    2.32%  \n",
       "train - fold4     5.76%      22.20%   0.68%    2.00%  \n",
       "val - fold4       4.15%      23.11%   0.73%    2.11%  "
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check distribution\n",
    "import pandas as pd\n",
    "from collections import Counter\n",
    "\n",
    "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)]\n",
    "\n",
    "distrs = [get_distribution(y)]\n",
    "index = ['training set']\n",
    "\n",
    "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}')\n",
    "    index.append(f'val - fold{fold_ind}')\n",
    "    \n",
    "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": 56,
   "id": "7387470e-1ca7-4cfe-8163-fcdef8fcb8f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.719999999999998 3.7200000000000024 6.199999999999997 2.4499999999999993 6.190000000000001\n"
     ]
    }
   ],
   "source": [
    "train_val_df = pd.DataFrame(distrs, index=index, columns = [categories[i] for i in range(np.max(y) + 1)])\n",
    "fold0_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[1].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[2].to_list())))))\n",
    "fold1_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[3].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[4].to_list())))))\n",
    "fold2_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[5].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[6].to_list())))))\n",
    "fold3_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[7].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[8].to_list())))))\n",
    "fold4_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[9].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[10].to_list())))))\n",
    "print(fold0_gap,fold1_gap,fold2_gap,fold3_gap,fold4_gap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "32f91e65-a0fe-4c83-bf49-4b75a5f72da7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.7899999999999974 2.9000000000000017 4.98 1.9499999999999997 5.0399999999999965\n"
     ]
    }
   ],
   "source": [
    "fold0_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[2].to_list())))))\n",
    "fold1_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[4].to_list())))))\n",
    "fold2_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[6].to_list())))))\n",
    "fold3_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[8].to_list())))))\n",
    "fold4_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[10].to_list())))))\n",
    "print(fold0_gap_trainingset,fold1_gap_trainingset,fold2_gap_trainingset,fold3_gap_trainingset,fold4_gap_trainingset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2910f615-e476-4827-bda1-4aa2379debbb",
   "metadata": {},
   "source": [
    "- competition에서 주어진 training set과 test set이 random하게 split되었고, 두 set 간의 분포가 유사할 것이라는 가정하에, fold 내에서의 train, val의 차이가 가장 작으면서, valset이 기존 training set의 분포와 가장 비슷하기에 test set을 가장 잘 대변할 것으로 추측되는 fold3에서의 train과 val을 모델학습의 train과 val로 사용"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "c9018931-cf2e-4370-9197-f7734a4c5da4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[   1    1    1 ... 4882 4882 4882]\n",
      "[   0    3    3 ... 4881 4881 4881]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from collections import Counter\n",
    "\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(X,y, groups)):\n",
    "    \n",
    "    if fold_ind==3:\n",
    "        fold3_best_train = groups[train_idx]\n",
    "        fold3_best_val = groups[val_idx]\n",
    "    \n",
    "print(fold3_best_train)\n",
    "print(fold3_best_val)\n",
    "\n",
    "# train, val 각각으로 사용할 image index 추출"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "281baf8c-1271-435e-90de-dcdf795e7e39",
   "metadata": {},
   "source": [
    "- json으로 추출"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "38254d69-50cd-4f5d-9d03-657e54fe5706",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train length: 3902, val length: 981\n"
     ]
    }
   ],
   "source": [
    "train_idx_list = list(set(fold3_best_train))  # json파일 내의 images, annotations를 수정해야함\n",
    "val_idx_list = list(set(fold3_best_val))\n",
    "\n",
    "print(f'train length: {len(train_idx_list)}, val length: {len(val_idx_list)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "bd7f56c0-e61b-4ac4-910e-91109242c2f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3902 981\n"
     ]
    }
   ],
   "source": [
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "print(len(new_data_images_train), len(new_data_images_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "48db9141-a026-4558-a531-dff5b1641600",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18454 4690\n"
     ]
    }
   ],
   "source": [
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "print(len(new_ann_train), len(new_ann_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "f1b4ee3b-15f8-4eb4-8714-242ae160a5c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23144"
      ]
     },
     "execution_count": 125,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(json_file['annotations'])  # 18454 + 4690"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "b7461808-d2e1-438b-adc4-9dd5c7015fb6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3902 18454\n",
      "981 4690\n"
     ]
    }
   ],
   "source": [
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "print(len(json_train['images']), len(json_train['annotations']))\n",
    "print(len(json_val['images']), len(json_val['annotations']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "bbbf0335-93d9-48db-8b8d-f6fd78d210c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Json file로 생성\n",
    "\n",
    "import json\n",
    "\n",
    "with open('train_StfGKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('val_StfGKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d354f3e-67f4-44d3-bfb8-6a9dc1321031",
   "metadata": {},
   "source": [
    "## MultiLabel "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4114b8fb-93fe-4529-9506-61ddacc18da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from sklearn.model_selection import StratifiedGroupKFold\n",
    "\n",
    "# load json \n",
    "annotation = 'dataset/train.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]\n",
    "\n",
    "X = np.ones((len(data['annotations']),1))\n",
    "y = np.array([v[1] for v in var])  # class\n",
    "groups = np.array([v[0] for v in var])  # image(group)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1d74fcbf-3acd-4dea-8690-1521aa631c46",
   "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>image_id</th>\n",
       "      <th>category_id</th>\n",
       "      <th>area</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>257301.66</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>10402.56</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>26259.36</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>69096.17</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>24164.58</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id  category_id       area  id\n",
       "0         0            0  257301.66   0\n",
       "1         1            3   10402.56   1\n",
       "2         1            7   26259.36   2\n",
       "3         1            4   69096.17   3\n",
       "4         1            5   24164.58   4"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "train_df = pd.DataFrame(data['annotations'])\n",
    "train_df = train_df[['image_id', 'category_id', 'area','id']]\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "bb6410c2-e514-4d58-87c4-99bcf404acbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "count     23144.00000\n",
       "mean      92863.48116\n",
       "std      135979.14884\n",
       "min           0.00000\n",
       "25%        9997.00000\n",
       "50%       38938.00000\n",
       "75%      119122.00000\n",
       "max     1048371.00000\n",
       "Name: area, dtype: float64"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.options.display.float_format = '{:.5f}'.format\n",
    "train_df.area.apply(lambda x: int(x)).describe()  # bbox의 크기를 0~25, 25~50, 50~75, 75~100 총 4가지로 분류"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "abc7a4ae-94d6-4bfe-85d5-be47ecd87700",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4883 4883 4883\n"
     ]
    }
   ],
   "source": [
    "# 이미지 별 bbox 개수\n",
    "bbox_num = train_df.groupby('image_id')['id'].count().to_list()\n",
    "\n",
    "# 작은 box부터 큰 box까지 0,1,2,3, 이미지 별 bbox의 평균 크기 분류\n",
    "area_mean_df = train_df.groupby('image_id')['area'].mean()\n",
    "area_class = []  \n",
    "for i in range(4883):\n",
    "    if area_mean_df[i] <= 9997:\n",
    "        area_class.append(0)\n",
    "    elif area_mean_df[i] <= 38938:\n",
    "        area_class.append(1)\n",
    "    elif area_mean_df[i] <= 119122:\n",
    "        area_class.append(2)\n",
    "    else:\n",
    "        area_class.append(3)\n",
    "\n",
    "# 최빈 class \n",
    "category_group = train_df.groupby('image_id')['category_id'].value_counts()\n",
    "most_class_list = []\n",
    "for i in range(4883):\n",
    "    most_class = category_group[i].index.to_list()[0] # 동률일 경우 가장 빠른 index 선택\n",
    "    most_class_list.append(most_class)\n",
    "    \n",
    "print(len(bbox_num), len(area_class), len(most_class_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "25c94a8a-17cb-4f03-a580-a7cdc358f969",
   "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>image_id</th>\n",
       "      <th>bbox_num</th>\n",
       "      <th>area_class</th>\n",
       "      <th>most_class_list</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id  bbox_num  area_class  most_class_list\n",
       "0         0         1           3                0\n",
       "1         1         8           2                0\n",
       "2         2         1           3                3\n",
       "3         3         2           2                2\n",
       "4         4         6           1                1"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multi_label_df = pd.DataFrame({'image_id':[i for i in range(4883)], 'bbox_num':bbox_num, 'area_class':area_class, 'most_class_list':most_class_list})\n",
    "multi_label_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "02e4f481-130a-44f5-b1dc-caf028223bc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TRAIN: [array([ 1,  3,  4,  5,  7,  8,  9, 10, 11, 12])]\n",
      " TEST: [array([ 0,  2,  6, 16, 19, 23, 25, 28, 36, 38])]\n",
      "TRAIN: [array([ 0,  2,  3,  4,  6,  7,  8,  9, 10, 11])]\n",
      " TEST: [array([ 1,  5, 14, 20, 21, 31, 33, 34, 41, 42])]\n",
      "TRAIN: [array([ 0,  1,  2,  3,  4,  5,  6,  8,  9, 10])]\n",
      " TEST: [array([ 7, 13, 18, 22, 24, 30, 49, 50, 51, 53])]\n",
      "TRAIN: [array([ 0,  1,  2,  5,  6,  7,  9, 10, 12, 13])]\n",
      " TEST: [array([ 3,  4,  8, 11, 26, 27, 29, 32, 37, 40])]\n",
      "TRAIN: [array([ 0,  1,  2,  3,  4,  5,  6,  7,  8, 11])]\n",
      " TEST: [array([ 9, 10, 12, 15, 17, 35, 39, 54, 58, 66])]\n"
     ]
    }
   ],
   "source": [
    "from iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n",
    "\n",
    "cv = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=41)\n",
    "X = np.ones((len(bbox_num),1))\n",
    "\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(multi_label_df, multi_label_df[['bbox_num', 'area_class']])):\n",
    "    \n",
    "    print(\"TRAIN:\", [train_idx[:10]])\n",
    "    print(\" TEST:\", [val_idx[:10]])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "ff536747-7b55-423a-87ce-4252501de761",
   "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>General trash</th>\n",
       "      <th>Paper</th>\n",
       "      <th>Paper pack</th>\n",
       "      <th>Metal</th>\n",
       "      <th>Glass</th>\n",
       "      <th>Plastic</th>\n",
       "      <th>Styrofoam</th>\n",
       "      <th>Plastic bag</th>\n",
       "      <th>Battery</th>\n",
       "      <th>Clothing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>training set</th>\n",
       "      <td>17.14%</td>\n",
       "      <td>27.45%</td>\n",
       "      <td>3.88%</td>\n",
       "      <td>4.04%</td>\n",
       "      <td>4.24%</td>\n",
       "      <td>12.72%</td>\n",
       "      <td>5.46%</td>\n",
       "      <td>22.37%</td>\n",
       "      <td>0.69%</td>\n",
       "      <td>2.02%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold0</th>\n",
       "      <td>17.05%</td>\n",
       "      <td>27.49%</td>\n",
       "      <td>3.85%</td>\n",
       "      <td>4.12%</td>\n",
       "      <td>4.48%</td>\n",
       "      <td>12.48%</td>\n",
       "      <td>5.25%</td>\n",
       "      <td>22.60%</td>\n",
       "      <td>0.64%</td>\n",
       "      <td>2.05%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold0</th>\n",
       "      <td>17.48%</td>\n",
       "      <td>27.27%</td>\n",
       "      <td>3.99%</td>\n",
       "      <td>3.75%</td>\n",
       "      <td>3.30%</td>\n",
       "      <td>13.67%</td>\n",
       "      <td>6.29%</td>\n",
       "      <td>21.47%</td>\n",
       "      <td>0.86%</td>\n",
       "      <td>1.92%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold1</th>\n",
       "      <td>17.07%</td>\n",
       "      <td>27.44%</td>\n",
       "      <td>4.03%</td>\n",
       "      <td>4.18%</td>\n",
       "      <td>4.26%</td>\n",
       "      <td>12.72%</td>\n",
       "      <td>5.46%</td>\n",
       "      <td>22.09%</td>\n",
       "      <td>0.65%</td>\n",
       "      <td>2.09%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold1</th>\n",
       "      <td>17.38%</td>\n",
       "      <td>27.46%</td>\n",
       "      <td>3.28%</td>\n",
       "      <td>3.50%</td>\n",
       "      <td>4.20%</td>\n",
       "      <td>12.69%</td>\n",
       "      <td>5.45%</td>\n",
       "      <td>23.48%</td>\n",
       "      <td>0.83%</td>\n",
       "      <td>1.74%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold2</th>\n",
       "      <td>17.18%</td>\n",
       "      <td>27.08%</td>\n",
       "      <td>3.79%</td>\n",
       "      <td>3.83%</td>\n",
       "      <td>4.21%</td>\n",
       "      <td>13.18%</td>\n",
       "      <td>5.55%</td>\n",
       "      <td>22.56%</td>\n",
       "      <td>0.74%</td>\n",
       "      <td>1.89%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold2</th>\n",
       "      <td>16.96%</td>\n",
       "      <td>28.91%</td>\n",
       "      <td>4.23%</td>\n",
       "      <td>4.91%</td>\n",
       "      <td>4.38%</td>\n",
       "      <td>10.86%</td>\n",
       "      <td>5.10%</td>\n",
       "      <td>21.63%</td>\n",
       "      <td>0.48%</td>\n",
       "      <td>2.55%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold3</th>\n",
       "      <td>17.08%</td>\n",
       "      <td>27.86%</td>\n",
       "      <td>3.91%</td>\n",
       "      <td>4.18%</td>\n",
       "      <td>4.09%</td>\n",
       "      <td>12.37%</td>\n",
       "      <td>5.45%</td>\n",
       "      <td>22.14%</td>\n",
       "      <td>0.74%</td>\n",
       "      <td>2.17%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold3</th>\n",
       "      <td>17.34%</td>\n",
       "      <td>25.80%</td>\n",
       "      <td>3.73%</td>\n",
       "      <td>3.49%</td>\n",
       "      <td>4.86%</td>\n",
       "      <td>14.11%</td>\n",
       "      <td>5.48%</td>\n",
       "      <td>23.30%</td>\n",
       "      <td>0.46%</td>\n",
       "      <td>1.43%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>train - fold4</th>\n",
       "      <td>17.29%</td>\n",
       "      <td>27.36%</td>\n",
       "      <td>3.80%</td>\n",
       "      <td>3.91%</td>\n",
       "      <td>4.18%</td>\n",
       "      <td>12.84%</td>\n",
       "      <td>5.58%</td>\n",
       "      <td>22.48%</td>\n",
       "      <td>0.66%</td>\n",
       "      <td>1.91%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>val - fold4</th>\n",
       "      <td>16.51%</td>\n",
       "      <td>27.80%</td>\n",
       "      <td>4.16%</td>\n",
       "      <td>4.60%</td>\n",
       "      <td>4.49%</td>\n",
       "      <td>12.23%</td>\n",
       "      <td>4.95%</td>\n",
       "      <td>21.96%</td>\n",
       "      <td>0.81%</td>\n",
       "      <td>2.49%</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              General trash   Paper Paper pack  Metal  Glass Plastic  \\\n",
       "training set         17.14%  27.45%      3.88%  4.04%  4.24%  12.72%   \n",
       "train - fold0        17.05%  27.49%      3.85%  4.12%  4.48%  12.48%   \n",
       "val - fold0          17.48%  27.27%      3.99%  3.75%  3.30%  13.67%   \n",
       "train - fold1        17.07%  27.44%      4.03%  4.18%  4.26%  12.72%   \n",
       "val - fold1          17.38%  27.46%      3.28%  3.50%  4.20%  12.69%   \n",
       "train - fold2        17.18%  27.08%      3.79%  3.83%  4.21%  13.18%   \n",
       "val - fold2          16.96%  28.91%      4.23%  4.91%  4.38%  10.86%   \n",
       "train - fold3        17.08%  27.86%      3.91%  4.18%  4.09%  12.37%   \n",
       "val - fold3          17.34%  25.80%      3.73%  3.49%  4.86%  14.11%   \n",
       "train - fold4        17.29%  27.36%      3.80%  3.91%  4.18%  12.84%   \n",
       "val - fold4          16.51%  27.80%      4.16%  4.60%  4.49%  12.23%   \n",
       "\n",
       "              Styrofoam Plastic bag Battery Clothing  \n",
       "training set      5.46%      22.37%   0.69%    2.02%  \n",
       "train - fold0     5.25%      22.60%   0.64%    2.05%  \n",
       "val - fold0       6.29%      21.47%   0.86%    1.92%  \n",
       "train - fold1     5.46%      22.09%   0.65%    2.09%  \n",
       "val - fold1       5.45%      23.48%   0.83%    1.74%  \n",
       "train - fold2     5.55%      22.56%   0.74%    1.89%  \n",
       "val - fold2       5.10%      21.63%   0.48%    2.55%  \n",
       "train - fold3     5.45%      22.14%   0.74%    2.17%  \n",
       "val - fold3       5.48%      23.30%   0.46%    1.43%  \n",
       "train - fold4     5.58%      22.48%   0.66%    1.91%  \n",
       "val - fold4       4.95%      21.96%   0.81%    2.49%  "
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "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)]\n",
    "\n",
    "var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]\n",
    "\n",
    "X = np.ones((len(data['annotations']),1))\n",
    "y = np.array([v[1] for v in var])  # class\n",
    "groups = np.array([v[0] for v in var])  # image(group)\n",
    "groups\n",
    "distrs = [get_distribution(y)]\n",
    "index = ['training set']\n",
    "origin_ann_df = pd.DataFrame(data['annotations'])\n",
    "\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(multi_label_df, multi_label_df[['bbox_num', 'area_class']])):\n",
    "    # train_y, val_y = y[train_idx], y[val_idx]\n",
    "    # train_gr, val_gr = groups[train_idx], groups[val_idx]\n",
    "    \n",
    "    train_y = origin_ann_df[origin_ann_df['image_id'].isin(train_idx)].category_id.to_list()\n",
    "    val_y = origin_ann_df[origin_ann_df['image_id'].isin(val_idx)].category_id.to_list()\n",
    "    # train_gr, val_gr = train_idx, val_idx\n",
    "    \n",
    "    \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}')\n",
    "    index.append(f'val - fold{fold_ind}')\n",
    "    \n",
    "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": 194,
   "id": "4bd2a081-199f-42b1-ae51-305fde9ce033",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.050000000000001 3.7799999999999994 8.360000000000001 7.910000000000001 5.0699999999999985\n"
     ]
    }
   ],
   "source": [
    "train_val_df = pd.DataFrame(distrs, index=index, columns = [categories[i] for i in range(np.max(y) + 1)])\n",
    "fold0_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[1].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[2].to_list())))))\n",
    "fold1_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[3].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[4].to_list())))))\n",
    "fold2_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[5].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[6].to_list())))))\n",
    "fold3_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[7].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[8].to_list())))))\n",
    "fold4_gap = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[9].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[10].to_list())))))\n",
    "print(fold0_gap,fold1_gap,fold2_gap,fold3_gap,fold4_gap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "160e558c-c609-45f2-9784-ec76d2f1e605",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.810000000000002 3.000000000000001 6.700000000000005 6.329999999999996 4.07\n"
     ]
    }
   ],
   "source": [
    "fold0_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[2].to_list())))))\n",
    "fold1_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[4].to_list())))))\n",
    "fold2_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[6].to_list())))))\n",
    "fold3_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[8].to_list())))))\n",
    "fold4_gap_trainingset = np.sum(np.abs(np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[0].to_list()))) - np.array(list(map(lambda x:float(x[:-1]), train_val_df.iloc[10].to_list())))))\n",
    "print(fold0_gap_trainingset,fold1_gap_trainingset,fold2_gap_trainingset,fold3_gap_trainingset,fold4_gap_trainingset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45e1462c-c177-48fe-aecf-4073fb39acac",
   "metadata": {},
   "source": [
    "- competition에서 주어진 training set과 test set이 random하게 split되었고, 두 set 간의 분포가 유사할 것이라는 가정하에, fold 내에서의 train, val의 차이가 가장 작으면서, valset이 기존 training set의 분포와 가장 비슷하기에 test set을 가장 잘 대변할 것으로 추측되는 fold2에서의 train과 val을 모델학습의 train과 val로 사용"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3440bf58-b8f4-4882-8eb5-b21d00e0863e",
   "metadata": {},
   "source": [
    "- 'most_class_list'는 빼고 multilabel stratified kfold를 진행했을 때, 포함시켰을때보다, 기존 training set()과의 누적분포차이가 약 0.2%정도 더 작았기에 , 'most_class_list'는 빼고 진행"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 278,
   "id": "6bbc7405-d5d9-4416-9639-b0a3e972e848",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0  2  3  4  6  7  8  9 10 11]\n",
      "[ 1  5 14 20 21 31 33 34 41 42]\n",
      "train length: 3907, val length: 976\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(multi_label_df, multi_label_df[['bbox_num', 'area_class']])):\n",
    "    \n",
    "    if fold_ind==1:\n",
    "        fold2_best_train = train_idx\n",
    "        fold2_best_val = val_idx\n",
    "    \n",
    "print(fold2_best_train[:10])\n",
    "print(fold2_best_val[:10])\n",
    "train_idx_list = fold2_best_train\n",
    "val_idx_list = fold2_best_val\n",
    "print(f'train length: {len(train_idx_list)}, val length: {len(val_idx_list)}')\n",
    "# train, val 각각으로 사용할 image index 추출"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "id": "67e066d5-0872-49a4-8480-2fdada889f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_train_df = multi_label_df[multi_label_df['image_id'].isin(train_idx_list)]\n",
    "multi_val_df = multi_label_df[multi_label_df['image_id'].isin(val_idx_list)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "id": "8d94637b-15f0-44e3-a095-174cc62e79ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train [0.405, 0.161, 0.083, 0.053, 0.045, 0.034, 0.033, 0.021, 0.022, 0.018, 0.013, 0.013, 0.01, 0.01, 0.009, 0.011, 0.009, 0.005, 0.005, 0.004, 0.036] \n",
      "\n",
      "val [0.409, 0.157, 0.082, 0.052, 0.044, 0.036, 0.031, 0.022, 0.02, 0.016, 0.015, 0.016, 0.008, 0.005, 0.008, 0.007, 0.006, 0.007, 0.005, 0.006, 0.046]\n"
     ]
    }
   ],
   "source": [
    "train_bbox_num_list = multi_train_df.bbox_num.value_counts().sort_index().to_list()\n",
    "train_bbox_num = train_bbox_num_list[:20]\n",
    "train_bbox_num.append(sum(train_bbox_num_list[20:]))  # box 20개 이상인 경우는 하나로 묶음\n",
    "train_bbox_num_ratio = [round(i/3907,3) for i in train_bbox_num] # 마지막 값은 bbox 20개 이상인 이미지 개수, 그 전까지는 박스가 1 2 3 .. 19 20개인 이미지 개수\n",
    "\n",
    "val_bbox_num_list = multi_val_df.bbox_num.value_counts().sort_index().to_list()\n",
    "val_bbox_num = val_bbox_num_list[:20]\n",
    "val_bbox_num.append(sum(val_bbox_num_list[20:]))  # box 20개 이상인 경우는 하나로 묶음\n",
    "val_bbox_num_ratio = [round(i/976,3) for i in val_bbox_num]  # 마지막 값은 bbox 20개 이상인 이미지 개수, 그 전까지는 박스가 1 2 3 .. 19 20개인 이미지 개수\n",
    "\n",
    "print('train',train_bbox_num_ratio,'\\n')\n",
    "print('val',val_bbox_num_ratio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 276,
   "id": "8a6d1f79-0b8c-4825-aa0f-14cdd4e2ba66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train [0.005, 0.103, 0.391, 0.501] \n",
      "\n",
      "val [0.005, 0.115, 0.377, 0.503]\n"
     ]
    }
   ],
   "source": [
    "train_area_class_list = multi_train_df.area_class.value_counts().sort_index().to_list()\n",
    "train_area_class_ratio = [round(i/3907,3) for i in train_area_class_list] # 마지막 값은 bbox 20개 이상인 이미지 개수, 그 전까지는 박스가 1 2 3 .. 19 20개인 이미지 개수\n",
    "\n",
    "val_area_class_list = multi_val_df.area_class.value_counts().sort_index().to_list()\n",
    "val_area_class_ratio = [round(i/976,3) for i in val_area_class_num_list]  # 마지막 값은 bbox 20개 이상인 이미지 개수, 그 전까지는 박스가 1 2 3 .. 19 20개인 이미지 개수\n",
    "\n",
    "print('train',train_area_class_ratio,'\\n')\n",
    "print('val',val_area_class_ratio)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd74ea5-2770-47aa-889e-e564226e3224",
   "metadata": {},
   "source": [
    "- train과 val의 box의 class 분포, 이미지 내 box의 평균 크기 분포, 이미지 내 box의 개수 분포 모두 비슷함을 알 수 있다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7338c97e-15eb-4cae-ba02-5675af138482",
   "metadata": {},
   "source": [
    "- json으로 추출"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "6046bfbe-a579-46c3-8c99-26a3e005d4b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3907 976\n"
     ]
    }
   ],
   "source": [
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "print(len(new_data_images_train), len(new_data_images_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 280,
   "id": "50757160-6024-4711-b050-fcbc2162380f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18425 4719\n"
     ]
    }
   ],
   "source": [
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "print(len(new_ann_train), len(new_ann_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 283,
   "id": "0bd0dd98-f444-44b0-a795-4d5f90d7431c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23144"
      ]
     },
     "execution_count": 283,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(json_file['annotations'])  # 18425 + 4719"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 284,
   "id": "a39481a4-b530-42be-bdf0-daf1b24412a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3907 18425\n",
      "976 4719\n"
     ]
    }
   ],
   "source": [
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "print(len(json_train['images']), len(json_train['annotations']))\n",
    "print(len(json_val['images']), len(json_val['annotations']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 285,
   "id": "3916fb24-d3f1-4310-9b20-444ab7d75d6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Json file로 생성\n",
    "\n",
    "import json\n",
    "\n",
    "with open('train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b98d39-9d60-43ea-9c51-17cbf8ffd738",
   "metadata": {},
   "source": [
    "### 하위 33% 크기의 bbox를 가지고 있는 image 추출\n",
    "\n",
    "- 작은 bbox에 대한 mAP가 매우 낮기에, 작은 bbox에 대해 fine tuning 위해"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "46b5b531-baec-46ec-8e7f-f7c06f0f2b14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "annotation = 'dataset/train_MultiStfKFold.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    multi_json = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "166669ce-e21a-4a8b-9552-e4d16efeb4c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "count     18425.00000\n",
       "mean      93327.34619\n",
       "std      136631.82471\n",
       "min           0.56000\n",
       "25%       10063.68000\n",
       "50%       39062.06000\n",
       "75%      119405.10000\n",
       "max     1048371.21000\n",
       "Name: area, dtype: float64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(multi_json['annotations'])\n",
    "pd.options.display.float_format = '{:.5f}'.format\n",
    "df.area.describe()  # 원래 분포(train+val)의 분포와 매우 유사"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "475e425c-2e66-4ff9-8656-3ffdd6477c53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "index   12889.00000\n",
       "area    16305.26000\n",
       "Name: 6080, dtype: float64"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.area.sort_values().reset_index().loc[int(18425*0.33)]  # 하위 33%에 해당하는 bbox 크기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ae52f470-a49a-45a2-b163-2f51bb491055",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1469"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_box_df = df[df['area'] < 16306]\n",
    "small_box_df.image_id.nunique()\n",
    "# 하위 33% 크기의 bbox를 포함하는 image의 수, 기존 train 이미지의 약 37%에 해당"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ee31297d-1ae2-43f9-a650-bb79cc7b975c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1469"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_box_image_id_list = list(set(small_box_df.image_id.to_list()))\n",
    "small_box_images = [multi_json['images'][i] for i in range(len(multi_json['images'])) if multi_json['images'][i]['id'] in small_box_image_id_list]\n",
    "len(small_box_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "aa621590-005e-4f3a-89e4-fc324533ad87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14144"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_box_annotations = [multi_json['annotations'][i] for i in range(len(multi_json['annotations'])) if multi_json['annotations'][i]['image_id'] in small_box_image_id_list]\n",
    "len(small_box_annotations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "90da468d-6a48-42b9-aa00-c56c82e6f388",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1469 14144\n"
     ]
    }
   ],
   "source": [
    "multi_json_smallbox_train = multi_json.copy()\n",
    "\n",
    "multi_json_smallbox_train['images'] = small_box_images\n",
    "multi_json_smallbox_train['annotations'] = small_box_annotations\n",
    "\n",
    "print(len(multi_json_smallbox_train['images']), len(multi_json_smallbox_train['annotations']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "af593f85-a4fc-40b4-a427-e599c49ae9a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## bbox크기 하위 3분의1에 해당하는 image의 Json file로 생성\n",
    "\n",
    "import json\n",
    "\n",
    "with open('smallbox_train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(multi_json_smallbox_train, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69c42700-918a-47a7-b2e3-170db8c7ff8e",
   "metadata": {},
   "source": [
    "### 5fold OOF를 위한 json 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "4d74e944-3d0d-458e-8b94-c22d146d2e89",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from sklearn.model_selection import StratifiedGroupKFold\n",
    "\n",
    "# load json \n",
    "annotation = 'dataset/train.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]\n",
    "\n",
    "X = np.ones((len(data['annotations']),1))\n",
    "y = np.array([v[1] for v in var])  # class\n",
    "groups = np.array([v[0] for v in var])  # image(group)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "fc58dba2-fff2-48c0-ab7e-7a1467a29582",
   "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>image_id</th>\n",
       "      <th>category_id</th>\n",
       "      <th>area</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>257301.66000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>10402.56000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>26259.36000</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>69096.17000</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>24164.58000</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id  category_id         area  id\n",
       "0         0            0 257301.66000   0\n",
       "1         1            3  10402.56000   1\n",
       "2         1            7  26259.36000   2\n",
       "3         1            4  69096.17000   3\n",
       "4         1            5  24164.58000   4"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "train_df = pd.DataFrame(data['annotations'])\n",
    "train_df = train_df[['image_id', 'category_id', 'area','id']]\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d392bdb6-3dba-46f4-acf4-ec266f7416f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "count     23144.00000\n",
       "mean      92863.48116\n",
       "std      135979.14884\n",
       "min           0.00000\n",
       "25%        9997.00000\n",
       "50%       38938.00000\n",
       "75%      119122.00000\n",
       "max     1048371.00000\n",
       "Name: area, dtype: float64"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pd.options.display.float_format = '{:.5f}'.format\n",
    "# train_df.area.apply(lambda x: int(x)).describe()  # bbox의 크기를 0~25, 25~50, 50~75, 75~100 총 4가지로 분류"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "b598f709-95ce-42e3-b646-42e0dbed8b11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4883 4883 4883\n"
     ]
    }
   ],
   "source": [
    "# 이미지 별 bbox 개수\n",
    "bbox_num = train_df.groupby('image_id')['id'].count().to_list()\n",
    "\n",
    "# 작은 box부터 큰 box까지 0,1,2,3, 이미지 별 bbox의 평균 크기 분류\n",
    "area_mean_df = train_df.groupby('image_id')['area'].mean()\n",
    "area_class = []  \n",
    "for i in range(4883):\n",
    "    if area_mean_df[i] <= 9997:\n",
    "        area_class.append(0)\n",
    "    elif area_mean_df[i] <= 38938:\n",
    "        area_class.append(1)\n",
    "    elif area_mean_df[i] <= 119122:\n",
    "        area_class.append(2)\n",
    "    else:\n",
    "        area_class.append(3)\n",
    "\n",
    "# 최빈 class \n",
    "category_group = train_df.groupby('image_id')['category_id'].value_counts()\n",
    "most_class_list = []\n",
    "for i in range(4883):\n",
    "    most_class = category_group[i].index.to_list()[0] # 동률일 경우 가장 빠른 index 선택\n",
    "    most_class_list.append(most_class)\n",
    "    \n",
    "print(len(bbox_num), len(area_class), len(most_class_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "c7a81204-7ab9-4ee3-a544-a9fdc9afb3bd",
   "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>image_id</th>\n",
       "      <th>bbox_num</th>\n",
       "      <th>area_class</th>\n",
       "      <th>most_class_list</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id  bbox_num  area_class  most_class_list\n",
       "0         0         1           3                0\n",
       "1         1         8           2                0\n",
       "2         2         1           3                3\n",
       "3         3         2           2                2\n",
       "4         4         6           1                1"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multi_label_df = pd.DataFrame({'image_id':[i for i in range(4883)], 'bbox_num':bbox_num, 'area_class':area_class, 'most_class_list':most_class_list})\n",
    "multi_label_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "1d683c1a-e6c4-4684-9d7f-aed8cc076911",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TRAIN: [array([ 1,  3,  4,  5,  7,  8,  9, 10, 11, 12])]\n",
      " TEST: [array([ 0,  2,  6, 16, 19, 23, 25, 28, 36, 38])]\n",
      "TRAIN: [array([ 0,  2,  3,  4,  6,  7,  8,  9, 10, 11])]\n",
      " TEST: [array([ 1,  5, 14, 20, 21, 31, 33, 34, 41, 42])]\n",
      "TRAIN: [array([ 0,  1,  2,  3,  4,  5,  6,  8,  9, 10])]\n",
      " TEST: [array([ 7, 13, 18, 22, 24, 30, 49, 50, 51, 53])]\n",
      "TRAIN: [array([ 0,  1,  2,  5,  6,  7,  9, 10, 12, 13])]\n",
      " TEST: [array([ 3,  4,  8, 11, 26, 27, 29, 32, 37, 40])]\n",
      "TRAIN: [array([ 0,  1,  2,  3,  4,  5,  6,  7,  8, 11])]\n",
      " TEST: [array([ 9, 10, 12, 15, 17, 35, 39, 54, 58, 66])]\n"
     ]
    }
   ],
   "source": [
    "from iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n",
    "\n",
    "cv = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=41)\n",
    "X = np.ones((len(bbox_num),1))\n",
    "\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(multi_label_df, multi_label_df[['bbox_num', 'area_class']])):\n",
    "    \n",
    "    print(\"TRAIN:\", [train_idx[:10]])\n",
    "    print(\" TEST:\", [val_idx[:10]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1a4da398-7660-498f-81d1-973b561d61d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "all_fold_train_list = []\n",
    "all_fold_val_list = []\n",
    "for fold_ind, (train_idx, val_idx) in enumerate(cv.split(multi_label_df, multi_label_df[['bbox_num', 'area_class']])):\n",
    "    \n",
    "    all_fold_train_list.append(list(train_idx))\n",
    "    all_fold_val_list.append(list(val_idx))\n",
    "    \n",
    "print(len(all_fold_train_list))\n",
    "print(len(all_fold_val_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "f36d660b-e7e8-45e7-89ea-781741630aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_idx_list = all_fold_train_list[0]\n",
    "val_idx_list = all_fold_val_list[0]\n",
    "\n",
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "\n",
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "with open('1___train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('1___val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "aa045ea0-3cd3-4302-a8c5-b888ec8ab336",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_idx_list = all_fold_train_list[1]\n",
    "val_idx_list = all_fold_val_list[1]\n",
    "\n",
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "\n",
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "with open('2___train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('2___val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d756bf95-1307-4402-83e7-42a2fd47a54a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_idx_list = all_fold_train_list[2]\n",
    "val_idx_list = all_fold_val_list[2]\n",
    "\n",
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "\n",
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "with open('3___train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('3___val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "6710cdb0-33da-4f80-8355-6abc935962f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_idx_list = all_fold_train_list[3]\n",
    "val_idx_list = all_fold_val_list[3]\n",
    "\n",
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "\n",
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "with open('4___train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('4___val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "c92efbcc-8af1-4100-8a81-5b797266d75d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_idx_list = all_fold_train_list[4]\n",
    "val_idx_list = all_fold_val_list[4]\n",
    "\n",
    "json_file = data.copy()\n",
    "new_data_images_train = [json_file['images'][i] for i in train_idx_list]\n",
    "new_data_images_val = [json_file['images'][i] for i in val_idx_list]\n",
    "\n",
    "new_ann_train = []\n",
    "new_ann_val = []\n",
    "for ann_id in range(len(data['annotations'])):\n",
    "    ann_img_id = json_file['annotations'][ann_id]['image_id']\n",
    "    if ann_img_id in train_idx_list:\n",
    "        new_ann_train.append(json_file['annotations'][ann_id])\n",
    "    if ann_img_id in val_idx_list:\n",
    "        new_ann_val.append(json_file['annotations'][ann_id])\n",
    "        \n",
    "json_train = json_file.copy()\n",
    "json_val = json_file.copy()\n",
    "\n",
    "json_train['images'] = new_data_images_train\n",
    "json_train['annotations'] = new_ann_train\n",
    "\n",
    "json_val['images'] = new_data_images_val\n",
    "json_val['annotations'] = new_ann_val\n",
    "\n",
    "with open('5___train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(json_train, t)\n",
    "\n",
    "with open('5___val_MultiStfKFold.json', 'w') as v:\n",
    "    json.dump(json_val, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc25add-13ac-4114-bd16-17c09f2197be",
   "metadata": {},
   "source": [
    "### 가로, 세로 중, 한 쪽이 다른 쪽의 6배 이상 큰 경우는 noise로 분류\n",
    "### bbox area가 1000이하인 작은 bbox는 noise로 분류\n",
    "### -> 제거한 annotation json 생성"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "0486ea36-3e58-48db-aaa4-0559ba6d2923",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['info', 'licenses', 'images', 'categories', 'annotations'])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# load json \n",
    "annotation = './train_MultiStfKFold.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    multi_json = json.load(f)\n",
    "multi_json.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "3b721443-e4a8-4edd-97fb-ac486e1395ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000.56"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_over1000_list = [multi_json['annotations'][i] for i in range(len(multi_json['annotations'])) if multi_json['annotations'][i]['area'] > 1000]\n",
    "min([area_over1000_list[i]['area'] for i in range(len(area_over1000_list))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "d4fbfb4e-67ce-4520-bbcb-6808ff4d3670",
   "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>image_id</th>\n",
       "      <th>category_id</th>\n",
       "      <th>area</th>\n",
       "      <th>bbox</th>\n",
       "      <th>iscrowd</th>\n",
       "      <th>id</th>\n",
       "      <th>area_ratio</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>257301.66</td>\n",
       "      <td>[197.6, 193.7, 547.8, 469.7]</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1.166276</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>3</td>\n",
       "      <td>324010.80</td>\n",
       "      <td>[267.9, 165.2, 631.6, 513.0]</td>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "      <td>1.231189</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>59550.94</td>\n",
       "      <td>[462.2, 369.4, 233.9, 254.6]</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>0.918696</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>6</td>\n",
       "      <td>80710.56</td>\n",
       "      <td>[773.3, 3.0, 188.4, 428.4]</td>\n",
       "      <td>0</td>\n",
       "      <td>11</td>\n",
       "      <td>0.439776</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>14768.88</td>\n",
       "      <td>[567.5, 462.2, 165.2, 89.4]</td>\n",
       "      <td>0</td>\n",
       "      <td>12</td>\n",
       "      <td>1.847875</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id  category_id       area                          bbox  iscrowd  \\\n",
       "0         0            0  257301.66  [197.6, 193.7, 547.8, 469.7]        0   \n",
       "1         2            3  324010.80  [267.9, 165.2, 631.6, 513.0]        0   \n",
       "2         3            2   59550.94  [462.2, 369.4, 233.9, 254.6]        0   \n",
       "3         3            6   80710.56    [773.3, 3.0, 188.4, 428.4]        0   \n",
       "4         4            1   14768.88   [567.5, 462.2, 165.2, 89.4]        0   \n",
       "\n",
       "   id  area_ratio  \n",
       "0   0    1.166276  \n",
       "1   9    1.231189  \n",
       "2  10    0.918696  \n",
       "3  11    0.439776  \n",
       "4  12    1.847875  "
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "area_over1000_df = pd.DataFrame(area_over1000_list)\n",
    "area_ratio_list = []\n",
    "for i, row in area_over1000_df.iterrows():\n",
    "    area_ratio = row['bbox'][2] / row['bbox'][3]\n",
    "    area_ratio_list.append(area_ratio)\n",
    "area_over1000_df['area_ratio'] = area_ratio_list\n",
    "area_over1000_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "f2b32a74-d849-46e0-b2ce-69a40355e94b",
   "metadata": {},
   "outputs": [],
   "source": [
    "area_over1000_df = area_over1000_df[(area_over1000_df['area_ratio']<6) & (area_over1000_df['area_ratio']>0.17)]\n",
    "area_ratio_filtered_list = area_over1000_df.id.to_list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "074ca703-ca05-4e7a-b378-c1795f252a89",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.026919945725915877"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(18425-17929)/18425 # 약 2.6%의 bbox 제거됨"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "7d01866d-8974-4f38-b16d-a9b0c4a8f3db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "17929"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_ratio_filtered_ann = [multi_json['annotations'][i] for i in range(len(multi_json['annotations'])) if multi_json['annotations'][i]['id'] in area_ratio_filtered_list]\n",
    "len(area_ratio_filtered_ann)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "c3437b6a-7b8c-4ee7-962a-9040ca60f52d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1075, 1324, 1474, 1672, 1967, 2930, 3782]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_image_list = list(set(area_over1000_df.image_id.to_list()))\n",
    "no_box_image_list = []\n",
    "for i in range(len(multi_json['images'])):\n",
    "    if multi_json['images'][i]['id'] not in new_image_list:\n",
    "        no_box_image_list.append(multi_json['images'][i]['id'])\n",
    "no_box_image_list\n",
    "## bbox가 사라진 image 7장도 함께 제거"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "84d61b0b-8745-46c0-8f1a-e669de7800c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3900"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_ratio_filtered_image = [multi_json['images'][i] for i in range(len(multi_json['images'])) if multi_json['images'][i]['id'] not in no_box_image_list]\n",
    "len(area_ratio_filtered_image) ## 3900이면 됨"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "a41042b3-c08b-4975-a62a-57ae02a23201",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "area_ratio_filtered_json = multi_json.copy()\n",
    "\n",
    "area_ratio_filtered_json['images'] = area_ratio_filtered_image\n",
    "area_ratio_filtered_json['annotations'] = area_ratio_filtered_ann\n",
    "with open('AreaRatioFilter_train_MultiStfKFold.json', 'w') as t:\n",
    "    json.dump(area_ratio_filtered_json, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60d01db-e94e-420a-bb12-2b58d9364b65",
   "metadata": {},
   "source": [
    "### pseudo labeling 된 test data + 기존 trainset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "7f034a63-3039-48b3-ad77-d493eec0f485",
   "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>PredictionString</th>\n",
       "      <th>image_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7 0.9920914173126221 208.55441284179688 48.766...</td>\n",
       "      <td>test/0000.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5 0.9239843487739563 125.11742401123047 2.8685...</td>\n",
       "      <td>test/0001.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1 0.9707459211349487 77.59071350097656 270.659...</td>\n",
       "      <td>test/0002.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>9 0.8658463954925537 149.75205993652344 253.91...</td>\n",
       "      <td>test/0003.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1 0.81441330909729 191.17074584960938 237.0220...</td>\n",
       "      <td>test/0004.jpg</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                    PredictionString       image_id\n",
       "0  7 0.9920914173126221 208.55441284179688 48.766...  test/0000.jpg\n",
       "1  5 0.9239843487739563 125.11742401123047 2.8685...  test/0001.jpg\n",
       "2  1 0.9707459211349487 77.59071350097656 270.659...  test/0002.jpg\n",
       "3  9 0.8658463954925537 149.75205993652344 253.91...  test/0003.jpg\n",
       "4  1 0.81441330909729 191.17074584960938 237.0220...  test/0004.jpg"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "pseudo_df = pd.read_csv('output_best.csv')\n",
    "pseudo_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "715a36e9-94a8-4282-b76e-ea9e80f0f0cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['info', 'licenses', 'images', 'categories', 'annotations'])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# load json \n",
    "annotation = 'dataset/test.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    test = json.load(f)\n",
    "    \n",
    "test.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1cb4da1d-162a-4bee-8a77-f8e516f2ffa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation = 'dataset/train_MultiStfKFold.json'\n",
    "\n",
    "with open(annotation) as f:\n",
    "    train = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "db0cbdc6-2f24-46da-ab98-b339b6e2b397",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13114\n"
     ]
    }
   ],
   "source": [
    "ann_list = []\n",
    "bbox_id_cumulate = 23144\n",
    "for i, row in pseudo_df.iterrows():\n",
    "    pred_str = row['PredictionString']\n",
    "    pred_str_list = pred_str.split()\n",
    "    if len(pred_str_list)==0:\n",
    "        continue\n",
    "    for j in range(len(pred_str_list)//6):\n",
    "        if float(pred_str_list[j*6+1]) > 0.5:\n",
    "            image_id = 4883 + int(row['image_id'].split('.')[0][5:])\n",
    "            category_id = int(pred_str_list[j*6])\n",
    "            bbox = [round(float(pred_str_list[j*6+2]),1),round(float(pred_str_list[j*6+3]),1),\n",
    "                    round((float(pred_str_list[j*6+4]) - float(pred_str_list[j*6+2])),1),\n",
    "                    round((float(pred_str_list[j*6+5]) - float(pred_str_list[j*6+3])),1)]\n",
    "            \n",
    "            # for k in range(len(bbox)):\n",
    "            #     if k==2:\n",
    "            #         bbox[k] == 1024\n",
    "            #     if k==3:\n",
    "            #         bbox[k] == 0\n",
    "            \n",
    "            area = round(bbox[2]*bbox[3],2)\n",
    "            is_crowd = 0\n",
    "            bbox_id = bbox_id_cumulate\n",
    "            bbox_id_cumulate+=1\n",
    "            ann_dict = {\n",
    "                        'image_id': image_id,\n",
    "                        'category_id': category_id,\n",
    "                        'area': area,\n",
    "                        'bbox': bbox,\n",
    "                        'iscrowd':is_crowd,\n",
    "                        'id': bbox_id\n",
    "                        }\n",
    "            ann_list.append(ann_dict)\n",
    "        else:\n",
    "            continue\n",
    "print(len(ann_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "dd49ac74-4242-4f35-ac0e-cfa05d96f321",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4683\n"
     ]
    }
   ],
   "source": [
    "ann_df = pd.DataFrame(ann_list)\n",
    "pseudo_img_id_list = list(set(ann_df.image_id.to_list()))\n",
    "print(len(pseudo_img_id_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "4356b8b3-c6c2-46d7-b762-edd719bf69c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(test['images'])):\n",
    "    img_id = test['images'][i]['id']\n",
    "    img_id = img_id + 4883\n",
    "    test['images'][i]['id'] = img_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d0dc0b83-fc9b-4cfb-bef1-2801071f8f98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4683\n"
     ]
    }
   ],
   "source": [
    "pseudo_img_list = [test['images'][i] for i in range(len(test['images'])) if test['images'][i]['id'] in pseudo_img_id_list]\n",
    "print(len(pseudo_img_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "cb842348-959b-401d-9d3e-098f782d20a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "pseudo_json = train.copy() \n",
    "pseudo_json['images'] = train['images'] + pseudo_img_list\n",
    "pseudo_json['annotations'] = train['annotations'] + ann_list\n",
    "with open('pseudo+train_coco.json', 'w') as t:\n",
    "    json.dump(pseudo_json, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b242d82b-07bc-4c6a-b419-5b4030ec4e1e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "competition_effdet",
   "language": "python",
   "name": "competition_effdet"
  },
  "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.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}