ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from ensemble_boxes import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# csv 파일 로드\n",
    "file1 = './mmdetection/work_dirs/01_swin_bbox_safe/submission_best_bbox_mAP_epoch_16.csv'\n",
    "file2 = './mmdetection/work_dirs/cascade-mask-rcnn_convnext-t/submission_best_bbox_mAP_epoch_20.csv'\n",
    "file3 = './mmdetection/work_dirs/02_res2net101_bbox_safe/submission_best_bbox_mAP_epoch_16.csv'\n",
    "output_file = './results/ensemble.csv'\n",
    "\n",
    "df1 = pd.read_csv(file1)\n",
    "df2 = pd.read_csv(file2)\n",
    "df3 = pd.read_csv(file3)\n",
    "\n",
    "df1['source'] = file1\n",
    "df2['source'] = file2\n",
    "df3['source'] = file3\n",
    "\n",
    "\n",
    "the_weight = [0.4, 0.3, 0.3]    # wbf ensemble 비율 설정\n",
    "\n",
    "\n",
    "# 세 DataFrame을 하나로 결합\n",
    "df = pd.concat([df1, df2, df3])\n",
    "\n",
    "# 이미지의 너비와 높이\n",
    "IMAGE_WIDTH = 1920\n",
    "IMAGE_HEIGHT = 1080\n",
    "\n",
    "def convert_boxes(row):\n",
    "    return [(row['point1_x'] / IMAGE_WIDTH), (row['point1_y'] / IMAGE_HEIGHT), \n",
    "            (row['point3_x'] / IMAGE_WIDTH), (row['point3_y'] / IMAGE_HEIGHT)]\n",
    "\n",
    "\n",
    "df['boxes'] = df.apply(lambda row: convert_boxes(row), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 각 이미지에 대한 앙상블\n",
    "ensemble_results = []\n",
    "\n",
    "grouped = df.groupby('file_name')\n",
    "\n",
    "for name, group in grouped:\n",
    "    scores_list = []\n",
    "    boxes_list = []\n",
    "    labels_list = []\n",
    "    weights = []\n",
    "    \n",
    "    g2 = group.groupby('source')\n",
    "    for source, group2 in g2:\n",
    "        if source == file1:\n",
    "            idx = 0\n",
    "        elif source == file2:\n",
    "            idx = 1\n",
    "        else:\n",
    "            idx = 2\n",
    "        weights.append(the_weight[idx])\n",
    "\n",
    "        score_t = group2['confidence'].values.tolist()\n",
    "        box_t = group2['boxes'].values.tolist()\n",
    "        label_t = group2['class_id'].values.tolist()\n",
    "\n",
    "        scores_list.append(score_t)\n",
    "        boxes_list.append(box_t)\n",
    "        labels_list.append(label_t)\n",
    "    \n",
    "    \n",
    "\n",
    "    boxes, scores, labels = weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights = weights, iou_thr=0.5, skip_box_thr=0.049)\n",
    "    for box, score, label in zip(boxes, scores, labels):\n",
    "        ensemble_results.append({\n",
    "            'file_name': name,\n",
    "            'class_id': int(label),\n",
    "            'confidence': score,\n",
    "            'point1_x': box[0] * IMAGE_WIDTH,\n",
    "            'point1_y': box[1] * IMAGE_HEIGHT,\n",
    "            'point2_x': box[2] * IMAGE_WIDTH,\n",
    "            'point2_y': box[1] * IMAGE_HEIGHT,\n",
    "            'point3_x': box[2] * IMAGE_WIDTH,\n",
    "            'point3_y': box[3] * IMAGE_HEIGHT,\n",
    "            'point4_x': box[0] * IMAGE_WIDTH,\n",
    "            'point4_y': box[3] * IMAGE_HEIGHT,\n",
    "        })\n",
    "\n",
    "ensemble_df = pd.DataFrame(ensemble_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 앙상블 결과를 csv 파일로 저장\n",
    "ensemble_df.to_csv(output_file, index=False)\n",
    "print(len(ensemble_df))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openmmlab",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}