{"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":"## YOLOX dataset\n\nIn this notebook you can find code for splitting and transforming COTS dataset to YOLOX-type dataset. ","metadata":{}},{"cell_type":"code","source":"\"\"\"CLI for dataframe splitting on train and test\"\"\"\n\n\nfrom enum import Enum\nfrom pathlib import Path\nfrom typing import List, Tuple\n\nimport pandas as pd\n\n\nclass SplitType(Enum):\n    video: str = 'video'\n    length: str = 'length'\n\n\ndef _split_dataframe_by_videos(\n    dataframe: pd.DataFrame, train_video_ids: List[int], val_video_ids: List[int]\n) -> Tuple[pd.DataFrame, pd.DataFrame]:\n    train_dataframe = dataframe.loc[dataframe['video_id'].isin(train_video_ids)]\n    val_dataframe = dataframe.loc[dataframe['video_id'].isin(val_video_ids)]\n\n    return train_dataframe, val_dataframe\n\n\ndef _split_dataframe_by_length(\n    dataframe: pd.DataFrame, val_length: float = 0.2\n) -> Tuple[pd.DataFrame, pd.DataFrame]:\n    train_dataframes = []\n    val_dataframes = []\n\n    for unique_video_id in dataframe['video_id'].unique():\n        dataframe_part: pd.DataFrame = dataframe.loc[\n            dataframe['video_id'] == unique_video_id\n        ]\n\n        dataframe_part.sort_values(by='video_frame')\n\n        train_part = dataframe_part.iloc[: -int(val_length * dataframe_part.shape[0])]\n        val_part = dataframe_part.iloc[-int(val_length * dataframe_part.shape[0]) :]\n\n        train_dataframes.append(train_part)\n        val_dataframes.append(val_part)\n\n    train_df = pd.concat(train_dataframes)\n    val_df = pd.concat(val_dataframes)\n\n    return train_df, val_df\n\n\ndef split_dataframe(\n    original_dataframe_path: Path,\n    train_name: str = 'train_part.csv',\n    val_name: str = 'val_part.csv',\n    split_type: SplitType = SplitType.video,\n) -> None:\n    dataframe = pd.read_csv(original_dataframe_path)\n\n    if split_type == SplitType.video:\n        train_dataframe, val_dataframe = _split_dataframe_by_videos(\n            dataframe=dataframe, train_video_ids=[0, 1], val_video_ids=[2]\n        )\n    elif split_type == SplitType.length:\n        train_dataframe, val_dataframe = _split_dataframe_by_length(\n            dataframe=dataframe, val_length=0.3\n        )\n    else:\n        raise ValueError('No such dataset type')\n\n    train_dataframe.to_csv(train_name, index=False)\n    val_dataframe.to_csv(val_name, index=False)\n","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"Module with transforming of dataset to Coco format\"\"\"\n\nimport json\nimport os.path\nfrom pathlib import Path\nfrom typing import Any, Dict, List, Tuple, Union\n\nimport pandas as pd\nfrom tqdm import tqdm\n\n\nclass DatasetToCocoTransformer:\n    ANNOTATIONS_COLUMN = 'annotations'\n    VIDEO_ID_COLUMN = 'video_id'\n    VIDEO_FRAME_COLUMN = 'video_frame'\n\n    def __init__(\n        self,\n        annotation_dataframe: pd.DataFrame,\n        image_extension: str = '.jpg',\n        image_size: Tuple[int, int] = (1280, 720),\n        verbose: bool = True,\n    ) -> None:\n        self._annotation_dataframe = annotation_dataframe\n\n        self._image_extension = image_extension\n        self._image_size = image_size\n        self._verbose = verbose\n\n    def transform(self) -> Dict[str, Any]:\n        self._preprocess_dataframe()\n        width_image, height_image = self._image_size\n\n        images_info = []\n        annotations_info = []\n        categories_info = [{'supercategory': 'cot', 'id': 0, 'name': 'cot'}]\n        last_segmentation_id = -1\n\n        for index, curr_row in tqdm(\n            self._annotation_dataframe.iterrows(),\n            postfix='Transforming dataset...',\n            disable=not self._verbose,\n            total=self._annotation_dataframe.shape[0],\n        ):\n            annotations = curr_row[self.ANNOTATIONS_COLUMN]\n            video_id = curr_row[self.VIDEO_ID_COLUMN]\n            video_frame = curr_row[self.VIDEO_FRAME_COLUMN]\n\n            if len(annotations) > 0:\n                image_info = self._get_image_info(\n                    image_id=index,\n                    video_id=video_id,\n                    video_frame=video_frame,\n                    width_image=width_image,\n                    height_image=height_image,\n                )\n                (\n                    image_annotations_info,\n                    last_segmentation_id,\n                ) = self._get_image_annotations_info(\n                    image_id=index,\n                    last_segmentation_id=last_segmentation_id,\n                    coords=annotations,\n                )\n\n                images_info.append(image_info)\n                annotations_info.extend(image_annotations_info)\n\n        result = {\n            'images': images_info,\n            'annotations': annotations_info,\n            'categories': categories_info,\n        }\n\n        return result\n\n    @staticmethod\n    def save_json(obj: Any, path: Path) -> None:\n        with path.open(mode='w', encoding='UTF-8') as file:\n            json.dump(obj, fp=file)\n\n    def _get_image_info(\n        self,\n        image_id: int,\n        video_id: int,\n        video_frame: int,\n        width_image: int,\n        height_image: int,\n    ) -> Dict[str, Any]:\n        image_name = str(video_frame) + self._image_extension\n        image_rel_path = os.path.join(f'video_{str(video_id)}', image_name)\n\n        image_info = {\n            'file_name': image_rel_path,\n            'width': width_image,\n            'height': height_image,\n            'id': image_id,\n        }\n\n        return image_info\n\n    def _get_image_annotations_info(\n        self, image_id: int, last_segmentation_id: int, coords: List[Dict[str, int]]\n    ) -> Tuple[List[Dict[str, Any]], int]:\n        image_annotations = []\n\n        for curr_coords in coords:\n            transformed_coords = self._transform_bbox(bbox=curr_coords)\n\n            last_segmentation_id += 1\n            curr_image_annotation = {\n                'bbox': transformed_coords,\n                'category_id': 0,\n                'image_id': image_id,\n                'iscrowd': False,\n                'area': transformed_coords[2] * transformed_coords[3],\n                'id': last_segmentation_id,\n            }\n            image_annotations.append(curr_image_annotation)\n\n        return image_annotations, last_segmentation_id\n\n    def _preprocess_dataframe(self) -> None:\n        if isinstance(self._annotation_dataframe[self.ANNOTATIONS_COLUMN].loc[0], str):\n            self._annotation_dataframe[\n                self.ANNOTATIONS_COLUMN\n            ] = self._annotation_dataframe[self.ANNOTATIONS_COLUMN].apply(eval)\n\n        self._annotation_dataframe[self.VIDEO_ID_COLUMN] = self._annotation_dataframe[\n            self.VIDEO_ID_COLUMN\n        ].apply(int)\n        self._annotation_dataframe[\n            self.VIDEO_FRAME_COLUMN\n        ] = self._annotation_dataframe[self.VIDEO_FRAME_COLUMN].apply(int)\n\n    @staticmethod\n    def _transform_bbox(bbox: Dict[str, int]) -> List[float]:\n        x_min = int(bbox['x'])\n        y_min = int(bbox['y'])\n        width = int(bbox['width'])\n        height = int(bbox['height'])\n\n        return [x_min, y_min, width, height]\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-12-21T16:32:47.793868Z","iopub.execute_input":"2021-12-21T16:32:47.794378Z","iopub.status.idle":"2021-12-21T16:32:47.819015Z","shell.execute_reply.started":"2021-12-21T16:32:47.794345Z","shell.execute_reply":"2021-12-21T16:32:47.817945Z"},"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"Module with CLI for converting dataset to YoloV5 format\"\"\"\n\n\nfrom pathlib import Path\nfrom typing import Tuple\n\nimport pandas as pd\n\n\ndef transform_dataset_to_coco(\n    annotations_path: Path,\n    res_path: Path,\n    image_extension: str = '.jpg',\n    image_size: Tuple[int, int] = (1280, 720),\n    verbose: bool = True,\n) -> None:\n    dataframe = pd.read_csv(annotations_path)\n\n    dataset_to_yolo_transformer = DatasetToCocoTransformer(\n        annotation_dataframe=dataframe,\n        image_extension=image_extension,\n        image_size=image_size,\n        verbose=verbose,\n    )\n\n    result = dataset_to_yolo_transformer.transform()\n    dataset_to_yolo_transformer.save_json(result, res_path)\n","metadata":{"execution":{"iopub.status.busy":"2021-12-21T16:38:59.677833Z","iopub.execute_input":"2021-12-21T16:38:59.678351Z","iopub.status.idle":"2021-12-21T16:38:59.687590Z","shell.execute_reply.started":"2021-12-21T16:38:59.678292Z","shell.execute_reply":"2021-12-21T16:38:59.686375Z"},"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Let's split data","metadata":{}},{"cell_type":"markdown","source":"Using code above you can split data in two ways:\n\n1. By length. In current state last 30% of each video would be sent to validation.\n2. By video. In current state `video_0` and `video_1` would be used for trainig, `video_2` for validation.\n\nYou can change `SPLIT_TYPE` argument below to `SplitType.video`, to use second way.","metadata":{}},{"cell_type":"code","source":"DATAFRAME_PATH = Path('../input/tensorflow-great-barrier-reef/train.csv')\nSPLIT_TYPE = SplitType.length\n\nTRAIN_NAME = 'train_part.csv'\nVAL_NAME = 'val_part.csv'\n\nsplit_dataframe(\n    original_dataframe_path=DATAFRAME_PATH,\n    train_name=TRAIN_NAME,\n    val_name=VAL_NAME,\n    split_type=SPLIT_TYPE,\n)","metadata":{"execution":{"iopub.status.busy":"2021-12-21T16:39:52.820356Z","iopub.execute_input":"2021-12-21T16:39:52.821642Z","iopub.status.idle":"2021-12-21T16:39:52.983187Z","shell.execute_reply.started":"2021-12-21T16:39:52.821581Z","shell.execute_reply":"2021-12-21T16:39:52.981947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Let's transform data","metadata":{}},{"cell_type":"markdown","source":"YOLOX uses COCO-like dataset, so lets transform our splits to COCO annotations","metadata":{}},{"cell_type":"code","source":"TRAIN_RES_COCO_PATH = Path('train.json')\n\ntransform_dataset_to_coco(\n    annotations_path=Path(TRAIN_NAME),\n    res_path=Path(TRAIN_RES_COCO_PATH)\n)\n\nVAL_RES_COCO_PATH = Path('val.json')\n\ntransform_dataset_to_coco(\n    annotations_path=Path(VAL_NAME),\n    res_path=Path(VAL_RES_COCO_PATH)\n)","metadata":{"execution":{"iopub.status.busy":"2021-12-21T16:41:42.506008Z","iopub.execute_input":"2021-12-21T16:41:42.506814Z","iopub.status.idle":"2021-12-21T16:41:44.813977Z","shell.execute_reply.started":"2021-12-21T16:41:42.506754Z","shell.execute_reply":"2021-12-21T16:41:44.812945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now you have two files: `./train.json` and `./val.json`. You can use them to train your YOLOX!\n\nIf you've found a bug, let me know!","metadata":{}},{"cell_type":"markdown","source":"## If this notebook was helpful, please upvote!","metadata":{}}]}