{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Simply Convert Data to COCO Format\n\n- Split the training data to train and validation\n- Convert both to COCO Formatted JSON","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Some basic setup:\n# import some common libraries\nimport numpy as np\nimport pandas as pd\nimport os\nimport json\nimport random\nimport cv2\nimport matplotlib.pyplot as plt\nimport ast\nfrom tqdm import tqdm\nfrom sklearn.model_selection import train_test_split","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TOP_INPUT_DIR = \"/kaggle/input/global-wheat-detection/\"\nDATA_TRAIN_DIR = f\"{TOP_INPUT_DIR}train/\"\nDATA_TEST_DIR = f\"{TOP_INPUT_DIR}test/\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create COCO Base","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"coco_base = { \"info\": {},\n              \"licenses\": [], \n              \"images\": [],\n              \"annotations\": [],\n              \"categories\": []}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write dataset info to COCO Format ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base[\"info\"] = {\n    \"description\": \"Global Wheat Detection Dataset, Kaggle 2020\",\n    \"url\": \"https://www.kaggle.com/c/global-wheat-detection/data\",\n    \"version\": \"1.0\",\n    \"year\": 2020,\n    \"contributor\": \"http://www.global-wheat.com/contributors/\",\n    \"date_created\": \"2020/05/29\"\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write licenses to COCO Format","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base[\"licenses\"].append(\n    {\n        \"url\": \"https://opensource.org/licenses/MIT\",\n        \"id\": 1,\n        \"name\": \"MIT License\"\n    }\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write category to COCO Format","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base[\"categories\"].append({\"supercategory\": \"grain\",\"id\": 1,\"name\": \"wheat\"})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{TOP_INPUT_DIR}train.csv')\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"uniq_images = train_df.image_id.unique()\nlen(uniq_images)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### There are 3373 unique images in the training set","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"print(train_df.width.unique()[0])\nprint(train_df.height.unique()[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### All images are of size 1024x1024","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Group bboxes by image\n\nCreate one row per image - combining all bboxes for that image into one column. This helps significantly reduce the number of rows to traverse.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_bboxes_per_image(df):\n    \"\"\"author: @impiyush\"\"\"\n    # convert from string list, to python list\n    df.bbox = df.bbox.apply(ast.literal_eval)\n    # group-by on image id and return all bboxes for that image as a list of lists\n    return train_df.groupby('image_id')['bbox'].apply(list).reset_index(name='bboxes')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df_bboxes_grped = get_bboxes_per_image(train_df)\ntrain_df_bboxes_grped.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"assert train_df_bboxes_grped.shape[0] == len(uniq_images), \"Number of images differ when grouped\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Split to train and validation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train, df_valid = train_test_split(\n    train_df_bboxes_grped,\n    test_size=0.05,\n    random_state=32,\n    shuffle=True\n)\n\ndf_train = df_train.reset_index(drop=True)\ndf_valid = df_valid.reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(df_train.shape[0], df_valid.shape[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write images to COCO Format","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base_train = coco_base.copy()\ncoco_base_valid = coco_base.copy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_coco_images(df):\n    \"\"\"author: @impiyush\"\"\"\n    images = []\n    for _,img in enumerate(tqdm(df.image_id.unique())):\n        img_dict = {\"license\":1, \n                    \"height\":1024, \n                    \"width\":1024,\n                    \"id\":img}\n        img_dict[\"file_name\"] = f\"{img}.jpg\"\n        images.append(img_dict)\n\n    return images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base_train[\"images\"] = set_coco_images(df_train)\ncoco_base_valid[\"images\"] = set_coco_images(df_valid)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's check the first three images in the list","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base_train[\"images\"][:3]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base_valid[\"images\"][:3]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's also verify the number of images match, just for sanity","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"assert len(coco_base_train[\"images\"])==len(df_train), \"Number of images differ from df_train\"\nassert len(coco_base_valid[\"images\"])==len(df_valid), \"Number of images differ from df_valid\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Write annotations to COCO Format\n\nThe most important piece of this puzzle","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_coco_annotations(df):\n    \"\"\"author: @impiyush\"\"\"\n    annos = []\n    id_cnt = 1\n\n    for _,row in tqdm(df.iterrows(), total=len(df)):\n    #     print(row)\n        anno = {}\n        anno['segmentation'] = []\n        anno['iscrowd'] = 0\n        anno['image_id'] = row['image_id']\n        anno['category_id'] = 1\n        bboxes = row['bboxes']\n    #     annos.append(box)\n        for ix, box in enumerate(bboxes):\n            anno['bbox'] = box # x,y,w,h\n            anno['area'] = box[2] * box[3] # w*h\n            anno['id'] = f\"{id_cnt:05}\"\n            annos.append(anno.copy()) # copy is necessary here, otherwise it will always point to the last value of anno\n            id_cnt += 1\n    \n    return annos","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"coco_base_train['annotations'] = set_coco_annotations(df_train)\ncoco_base_valid['annotations'] = set_coco_annotations(df_valid)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dump COCO formatted JSON for train and validation sets","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"with open(f'/kaggle/working/coco_train.json','w') as train_coco:\n    json.dump(coco_base_train, train_coco)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with open(f'/kaggle/working/coco_valid.json','w') as valid_coco:\n    json.dump(coco_base_valid, valid_coco)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"### There you go, now you have your data in COCO Format. Hope you liked this quick Kernel!","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}