{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"<h1><center>Efficient YoloV5 Dataset Generator</center></h1>     \n\n<center><img src = \"https://i.imgur.com/iatgdo5.jpg\" width = \"635\" height = \"235\"/></center>         \n\nThis dataset was built to be compatible with the train (train.py) script that can be found [HERE](https://github.com/ultralytics/yolov5). I also have a training notebook with WandB integration that you can find [HERE](https://www.kaggle.com/coldfir3/yolov5-train/edit/run/81607643). The inference notebook is still WIP. The resulting kaggle Dataset cand be found [HERE](https://www.kaggle.com/coldfir3/great-barrier-reef-yolov5)\n\nThe tree main tasks into converting this dataset to Yolo format are:\n1. Splitting into train and test\n1. Converting the bboxes to yolo format `[xc, yc, w, h]` and saving them into text files\n1. Arranging the files in the expected folders and writting the `.yaml`\n\n<h3 style='background:orange; color:black'><center>Consider upvoting this notebook if you found it helpful.</center></h3>","metadata":{}},{"cell_type":"code","source":"import os\nfrom ast import literal_eval\n\nimport numpy as np\nimport pandas as pd\n\nfrom tqdm.notebook import tqdm\n\nimport shutil","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-12-06T13:27:46.888534Z","iopub.execute_input":"2021-12-06T13:27:46.888992Z","iopub.status.idle":"2021-12-06T13:27:46.986787Z","shell.execute_reply.started":"2021-12-06T13:27:46.888905Z","shell.execute_reply":"2021-12-06T13:27:46.985863Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_PATH = '/kaggle/input/tensorflow-great-barrier-reef/train_images'","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:27:46.988199Z","iopub.execute_input":"2021-12-06T13:27:46.988458Z","iopub.status.idle":"2021-12-06T13:27:46.992499Z","shell.execute_reply.started":"2021-12-06T13:27:46.988429Z","shell.execute_reply":"2021-12-06T13:27:46.991547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Reading the train data","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/tensorflow-great-barrier-reef/train.csv')\ndf.tail()","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:27:46.993986Z","iopub.execute_input":"2021-12-06T13:27:46.994440Z","iopub.status.idle":"2021-12-06T13:27:47.073984Z","shell.execute_reply.started":"2021-12-06T13:27:46.994398Z","shell.execute_reply":"2021-12-06T13:27:47.073389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Background images (no detections)\n\nFor this dataset I decided to use a total of `N=6000` images only. This is because most of the images don't have any annotation and Yolo recommends the folowing:\n> Background images. Background images are images with no objects that are added to a dataset to reduce False Positives (FP). We recommend about 0-10% background images to help reduce FPs (COCO has 1000 background images for reference, 1% of the total). No labels are required for background images.\n\nAs you can see below, the dataset have way more than that.","metadata":{}},{"cell_type":"code","source":"n_with_annotations = (df['annotations'] != '[]').sum()\nlen(df), n_with_annotations","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"N = 6000\ndf = pd.concat([\n    df[df['annotations'] != '[]'],\n    df[df['annotations'] == '[]'].sample(N - n_with_annotations)\n]).sample(frac=1).reset_index(drop = True)\ndf.tail()","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:29:33.564398Z","iopub.execute_input":"2021-12-06T13:29:33.564782Z","iopub.status.idle":"2021-12-06T13:29:33.583160Z","shell.execute_reply.started":"2021-12-06T13:29:33.564754Z","shell.execute_reply":"2021-12-06T13:29:33.582384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train Val split","metadata":{}},{"cell_type":"markdown","source":"For this data, I believe that the only consistent way to split the images between train/val are by video.","metadata":{}},{"cell_type":"code","source":"df['video_id'].value_counts()","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:30:59.652001Z","iopub.execute_input":"2021-12-06T13:30:59.652852Z","iopub.status.idle":"2021-12-06T13:30:59.660074Z","shell.execute_reply.started":"2021-12-06T13:30:59.652804Z","shell.execute_reply":"2021-12-06T13:30:59.659277Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"So we will use video 2 as validation and 0 and 1 as training. For better final performance I advise you to do a 3-fold split on this data and ensemble the final models using WBF that you can find [HERE](https://github.com/ZFTurbo/Weighted-Boxes-Fusion)","metadata":{}},{"cell_type":"code","source":"valid = df['video_id'] == 2\ntrain = df['video_id'] != 2\ndf.loc[valid, 'is_valid'] = True\ndf.loc[train, 'is_valid'] = False\n\ndf['annotations'] = df['annotations'].apply(literal_eval)\ndf['path'] = df.apply(lambda row: f\"{TRAIN_PATH}/video_{row['video_id']}/{row['video_frame']}.jpg\", axis = 1)\n\ndf.tail()","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:32:12.409915Z","iopub.execute_input":"2021-12-06T13:32:12.410256Z","iopub.status.idle":"2021-12-06T13:32:12.943588Z","shell.execute_reply.started":"2021-12-06T13:32:12.410202Z","shell.execute_reply":"2021-12-06T13:32:12.942743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The .yaml file\n\nYolo uses an .yaml file to indicate the number and name of the classe aswell as the location of the images/labels.","metadata":{}},{"cell_type":"code","source":"%%writefile config.yaml\n\ntrain: train \nval: valid\n\nnc: 1  \nnames: ['starfish'] ","metadata":{"execution":{"iopub.status.busy":"2021-12-06T13:32:23.133999Z","iopub.execute_input":"2021-12-06T13:32:23.134304Z","iopub.status.idle":"2021-12-06T13:32:23.140596Z","shell.execute_reply.started":"2021-12-06T13:32:23.134269Z","shell.execute_reply":"2021-12-06T13:32:23.139694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Converting to yolo bbox format\n\nYolo uses a bbox format of `[xc, yc, w, y]` therefore we need to adjust our dataset to reflect that notation.\n\n<center><img src = \"https://user-images.githubusercontent.com/26833433/91506361-c7965000-e886-11ea-8291-c72b98c25eec.jpg\" width = \"635\" height = \"235\"/></center>        ","metadata":{}},{"cell_type":"code","source":"def to_yolo(box, img_w = 1280, img_h = 720):\n    \n    w = box['width']\n    h = box['height']\n    xc = box['x'] + int(np.round(w/2))\n    yc = box['y'] + int(np.round(h/2))\n\n    return [xc/img_w, yc/img_h, w/img_w, h/img_h]","metadata":{"execution":{"iopub.status.busy":"2021-12-05T20:00:23.892295Z","iopub.execute_input":"2021-12-05T20:00:23.892967Z","iopub.status.idle":"2021-12-05T20:00:23.900777Z","shell.execute_reply.started":"2021-12-05T20:00:23.892916Z","shell.execute_reply":"2021-12-05T20:00:23.899743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Looping through the dataframe and saving the files","metadata":{}},{"cell_type":"code","source":"os.makedirs('train/images', exist_ok=True)\nos.makedirs('train/labels', exist_ok=True)\nos.makedirs('valid/images', exist_ok=True)\nos.makedirs('valid/labels', exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2021-12-05T20:00:33.310573Z","iopub.execute_input":"2021-12-05T20:00:33.311159Z","iopub.status.idle":"2021-12-05T20:00:33.317808Z","shell.execute_reply.started":"2021-12-05T20:00:33.311096Z","shell.execute_reply":"2021-12-05T20:00:33.316819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i, row in tqdm(df.iterrows(), total = len(df)):\n    \n    bboxes = row['annotations']\n    bboxes = [to_yolo(bbox) for bbox in bboxes]\n    \n    base_dir = 'valid' if row['is_valid'] else 'train'\n    fname = f\"{row['video_id']}_{row['video_frame']}\"\n    \n    with open(f'{base_dir}/labels/{fname}.txt', 'w+') as f:\n        for bbox in bboxes:\n            f.write('0 ' + ' '.join([str(round(b, 3)) for b in bbox]) + '\\n')\n    shutil.copyfile(row['path'], f\"{base_dir}/images/{fname}.jpg\")","metadata":{"execution":{"iopub.status.busy":"2021-12-05T20:00:36.282694Z","iopub.execute_input":"2021-12-05T20:00:36.283242Z","iopub.status.idle":"2021-12-05T20:01:51.243427Z","shell.execute_reply.started":"2021-12-05T20:00:36.283201Z","shell.execute_reply":"2021-12-05T20:01:51.242424Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Ziping the files so kaggle can assemble a dataset\n\nthe final dataset can be found [HERE](https://www.kaggle.com/coldfir3/great-barrier-reef-yolov5)","metadata":{}},{"cell_type":"code","source":"shutil.make_archive('valid', 'zip', 'valid')\nshutil.make_archive('train', 'zip', 'train')","metadata":{"execution":{"iopub.status.busy":"2021-12-05T20:01:51.245522Z","iopub.execute_input":"2021-12-05T20:01:51.245811Z","iopub.status.idle":"2021-12-05T20:04:19.116933Z","shell.execute_reply.started":"2021-12-05T20:01:51.245777Z","shell.execute_reply":"2021-12-05T20:04:19.115911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree('valid') \nshutil.rmtree('train') ","metadata":{"execution":{"iopub.status.busy":"2021-12-05T20:04:19.118334Z","iopub.execute_input":"2021-12-05T20:04:19.118568Z","iopub.status.idle":"2021-12-05T20:04:20.066054Z","shell.execute_reply.started":"2021-12-05T20:04:19.118538Z","shell.execute_reply":"2021-12-05T20:04:20.065423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}