{"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":"# # Pipeline process dataset (Fast and easy to implement)\nI share my work end to end process dataset based on pytorch + albumentation, this pipeline easy to add or remove any augmentation and fast process with dataloder (pytorch). My contribution to prepare dataset ( box augmentation + enhance contrast) for yolov5/yolor and easy to modify data for yolox.\n\nThank great notebook [ Underwater img Enhancement + EDA](https://www.kaggle.com/soumya9977/learning-to-sea-underwater-img-enhancement-eda) from [somuSan](https://www.kaggle.com/soumya9977)","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport cv2\nimport os\nimport albumentations as A\nimport glob\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom fastprogress.fastprogress import master_bar, progress_bar\n# from more_itertools import chunked\nimport multiprocessing as mp\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.utils.data as Data\nfrom tqdm.notebook import tqdm\ntqdm.pandas()\nimport ast","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:02:53.312944Z","iopub.execute_input":"2022-01-02T08:02:53.313228Z","iopub.status.idle":"2022-01-02T08:02:53.319126Z","shell.execute_reply.started":"2022-01-02T08:02:53.313198Z","shell.execute_reply":"2022-01-02T08:02:53.318478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    path_original = \"../input/tensorflow-great-barrier-reef/train.csv\"\n    fold_index = 0\n    path_dataset = \"./aug\"\n    os.makedirs(path_dataset, exist_ok=True)\n    visualize = False\n    worker=4\n    batch_size = 128\n    kfold = True\n    aug_box = True\n    aug_box_time = 8 ### total times augmentation\n    use_coco2yolo = False\n    ","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:02:53.336857Z","iopub.execute_input":"2022-01-02T08:02:53.337610Z","iopub.status.idle":"2022-01-02T08:02:53.357699Z","shell.execute_reply.started":"2022-01-02T08:02:53.337560Z","shell.execute_reply":"2022-01-02T08:02:53.357036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf $CFG.path_dataset","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:02:53.358995Z","iopub.execute_input":"2022-01-02T08:02:53.359775Z","iopub.status.idle":"2022-01-02T08:02:54.161454Z","shell.execute_reply.started":"2022-01-02T08:02:53.359736Z","shell.execute_reply":"2022-01-02T08:02:54.160264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_path(row):\n    row['image_path'] = f'../input/tensorflow-great-barrier-reef/train_images/video_{row.video_id}/{row.video_frame}.jpg'\n    return row","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:02:54.163211Z","iopub.execute_input":"2022-01-02T08:02:54.163493Z","iopub.status.idle":"2022-01-02T08:02:54.169377Z","shell.execute_reply.started":"2022-01-02T08:02:54.163457Z","shell.execute_reply":"2022-01-02T08:02:54.168237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.kfold:\n    from sklearn.model_selection import GroupKFold\n    df = pd.read_csv(CFG.path_original)\n    df = df.progress_apply(get_path, axis=1)\n    FDA_image = df[df['annotations']=='[]']['image_path'].tolist()\n    df = df[df['annotations'] != '[]']\n    df['annotations'] = df['annotations'].progress_apply(lambda x: ast.literal_eval(x))\n    \n    kf = GroupKFold(n_splits = 5)\n    df = df.reset_index(drop=True)\n    df['fold'] = -1\n    for fold, (train_idx, val_idx) in enumerate(kf.split(df, y = df.video_id.tolist(), groups=df.sequence)):\n        df.loc[val_idx, 'fold'] = fold\n    display(df.fold.value_counts())\ntrain = df[df['fold']!=CFG.fold_index]\nvalid = df[df['fold']==CFG.fold_index]","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:02:54.171283Z","iopub.execute_input":"2022-01-02T08:02:54.172153Z","iopub.status.idle":"2022-01-02T08:03:10.472571Z","shell.execute_reply.started":"2022-01-02T08:02:54.172109Z","shell.execute_reply":"2022-01-02T08:03:10.471988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_bbox(annots):\n    bboxes = [list(annot.values()) for annot in annots]\n    return bboxes","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.473487Z","iopub.execute_input":"2022-01-02T08:03:10.474173Z","iopub.status.idle":"2022-01-02T08:03:10.477928Z","shell.execute_reply.started":"2022-01-02T08:03:10.474141Z","shell.execute_reply":"2022-01-02T08:03:10.477397Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train['bboxes'] = train.annotations.progress_apply(get_bbox)\nvalid['bboxes'] = valid.annotations.progress_apply(get_bbox)","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.479147Z","iopub.execute_input":"2022-01-02T08:03:10.479448Z","iopub.status.idle":"2022-01-02T08:03:10.707715Z","shell.execute_reply.started":"2022-01-02T08:03:10.479422Z","shell.execute_reply":"2022-01-02T08:03:10.706730Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for folder_type in ['images', 'labels']:\n    path_phase = CFG.path_dataset + '/' + folder_type\n    os.makedirs(path_phase, exist_ok=True)\n    for phase in ['train', 'valid']:\n        path_type = path_phase + '/' + phase\n        os.makedirs(path_type, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.709231Z","iopub.execute_input":"2022-01-02T08:03:10.709606Z","iopub.status.idle":"2022-01-02T08:03:10.716434Z","shell.execute_reply.started":"2022-01-02T08:03:10.709561Z","shell.execute_reply":"2022-01-02T08:03:10.715577Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### aug image ###\ndef write_box_into_image(bouding_box, path_image):\n    ### bouding_box, path_image\n    image = cv2.imread(path_image)\n#     print(image.shape)\n    for bb in bouding_box:\n        x, y, w, h = bb['x'], bb['y'], bb['width'], bb['height']\n        print(bb)\n        image = cv2.rectangle(image, (x,y), (x+w,y+h), (0,0,255),2)\n    return image\nclass HE_HSV(A.ImageOnlyTransform):\n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, image,**params):\n        img_hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)\n\n        # Histogram equalisation on the V-channel\n        img_hsv[:, :, 2] = cv2.equalizeHist(img_hsv[:, :, 2])\n\n        # convert image back from HSV to RGB\n        image_hsv = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)\n\n        return image_hsv\n    \nclass RecoverHE(A.ImageOnlyTransform):\n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, sceneRadiance,**params):\n        for i in range(3):\n            sceneRadiance[:, :, i] =  cv2.equalizeHist(sceneRadiance[:, :, i])\n        return sceneRadiance\n\nclass CLAHE_HSV(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, img, **params):\n        hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n\n        h, s, v = hsv_img[:,:,0], hsv_img[:,:,1], hsv_img[:,:,2]\n        clahe = cv2.createCLAHE(clipLimit = 15.0, tileGridSize = (20,20))\n        v = clahe.apply(v)\n\n        hsv_img = np.dstack((h,s,v))\n\n        rgb = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB)\n\n        return rgb\n\nclass RecoverCLAHE(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, sceneRadiance, **params):\n        clahe = cv2.createCLAHE(clipLimit=7, tileGridSize=(14, 14))\n        for i in range(3):\n            sceneRadiance[:, :, i] = clahe.apply((sceneRadiance[:, :, i]))\n\n        return sceneRadiance\n\nclass Gamma_enhancement(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        self.gamma = 1/0.6\n        self.R = 255.0\n        \n    def apply(self, image, **params):\n        return (self.R * np.power(image.astype(np.uint32)/self.R, self.gamma)).astype(np.uint8)\n\nclass RecoverGC(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, sceneRadiance, **params):\n        sceneRadiance = sceneRadiance/255.0\n        # clahe = cv2.createCLAHE(clipLimit=2, tileGridSize=(2, 2))\n        for i in range(3):\n            sceneRadiance[:, :, i] =  np.power(sceneRadiance[:, :, i] / float(np.max(sceneRadiance[:, :, i])), 3.2)\n        sceneRadiance = np.clip(sceneRadiance*255, 0, 255)\n        sceneRadiance = np.uint8(sceneRadiance)\n        return sceneRadiance\n\nclass RecoverICM(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=False):\n        super().__init__(always_apply, p)\n        \n    def apply(self, image, **params):\n        img_stre = stretching(iamge)\n        sceneRadiance = sceneRadianceRGB(img_stre)\n        sceneRadiance = HSVStretching(sceneRadiance)\n        sceneRadiance = sceneRadianceRGB(sceneRadiance)\n\n        return sceneRadiance","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.717827Z","iopub.execute_input":"2022-01-02T08:03:10.718238Z","iopub.status.idle":"2022-01-02T08:03:10.737961Z","shell.execute_reply.started":"2022-01-02T08:03:10.718208Z","shell.execute_reply":"2022-01-02T08:03:10.737309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_one_image(image):\n    plt.figure(figsize=(20, 15))\n    plt.gcf().set_dpi(100)\n    plt.imshow(image)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.738835Z","iopub.execute_input":"2022-01-02T08:03:10.739050Z","iopub.status.idle":"2022-01-02T08:03:10.753541Z","shell.execute_reply.started":"2022-01-02T08:03:10.739024Z","shell.execute_reply":"2022-01-02T08:03:10.752755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train = cross_validation[cross_validation['fold']!=CFG.fold_index]\n# valid = cross_validation[cross_validation['fold']==CFG.fold_index]","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.756642Z","iopub.execute_input":"2022-01-02T08:03:10.756966Z","iopub.status.idle":"2022-01-02T08:03:10.764241Z","shell.execute_reply.started":"2022-01-02T08:03:10.756923Z","shell.execute_reply":"2022-01-02T08:03:10.763129Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bbox_to_txt(bboxes):\n    \"\"\"\n    Convert a list of bbox into a string in YOLO format (to write a file).\n    @bboxes : numpy array of bounding boxes \n    return : a string for each object in new line: <object-class> <x> <y> <width> <height>\n    \"\"\"\n    txt=''\n    for index,l in enumerate(bboxes):\n        l = [str(x) for x in l[:4]]\n        l = ' '.join(l)\n        txt +=  '0' +' ' + l + '\\n'\n    return txt","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.765967Z","iopub.execute_input":"2022-01-02T08:03:10.766324Z","iopub.status.idle":"2022-01-02T08:03:10.778160Z","shell.execute_reply.started":"2022-01-02T08:03:10.766283Z","shell.execute_reply":"2022-01-02T08:03:10.777446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_data(df, total_aug=10):\n    ### return aug df \n    df_aug = pd.DataFrame(np.repeat(df.values, total_aug, axis=0), columns=df.columns)\n    df_aug['aug_index'] = np.repeat(np.arange(1,total_aug+1).reshape(1,-1),df.shape[0], axis=0).reshape(-1)\n    df_aug = df_aug.sample(frac=1)\n    df_aug = df_aug.reset_index(drop=True)\n    return df_aug\ntrain = prepare_data(train, CFG.aug_box_time)\nvalid = prepare_data(valid,1)","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.779592Z","iopub.execute_input":"2022-01-02T08:03:10.779875Z","iopub.status.idle":"2022-01-02T08:03:10.815162Z","shell.execute_reply.started":"2022-01-02T08:03:10.779836Z","shell.execute_reply":"2022-01-02T08:03:10.814454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AUG_DATASET(Dataset):\n    \n    def __init__(self, df, mode, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.mode = mode\n        self.transform = transform\n        \n    def coco2yolo(self, bboxes, image_height=720, image_width=1280):\n        \"\"\"\n        coco => [xmin, ymin, w, h]\n        yolo => [xmid, ymid, w, h] (normalized)\n        \"\"\"\n        bboxes = np.array(bboxes)\n        bboxes = bboxes.copy().astype(float) # otherwise all value will be 0 as voc_pascal dtype is np.int\n\n        # normolizinig\n        bboxes[..., [0, 2]]= bboxes[..., [0, 2]]/ image_width\n        bboxes[..., [1, 3]]= bboxes[..., [1, 3]]/ image_height\n\n        # converstion (xmin, ymin) => (xmid, ymid)\n        bboxes[..., [0, 1]] = bboxes[..., [0, 1]] + bboxes[..., [2, 3]]/2\n\n        return bboxes\n    \n    def coord_to_box(self, bouding_box, image):\n        box_yolo_format = []\n        height, width = image.shape[0], image.shape[1]\n        \n        if CFG.use_coco2yolo:\n            box_yolo_format = self.coco2yolo(bouding_box)\n            box_yolo_format = np.clip(box_yolo_format,0,1)\n            label = np.repeat([0],box_yolo_format.shape[0]).reshape(-1,1)\n            box_yolo_format = np.append(box_yolo_format,label, axis=1)\n        else:\n            for bb in bouding_box:\n                label = [max(0,bb[0]), max(0,bb[1]), min(bb[0]+bb[2], 1280), min(720,bb[1]+bb[3]), '0']\n                bbox_albu = A.convert_bbox_to_albumentations(label, source_format='pascal_voc', rows=height, cols=width)\n                bbox_yolo = A.convert_bbox_from_albumentations(bbox_albu, target_format='yolo', rows=height, cols=width, check_validity=True)\n                clip_box = [np.clip(value,0,1) for value in bbox_yolo[:-1]] + [bbox_yolo[-1]]\n                box_yolo_format.append(clip_box)\n        return box_yolo_format\n\n    def bbox_to_txt(self, bboxes):\n        \"\"\"\n        Convert a list of bbox into a string in YOLO format (to write a file).\n        @bboxes : numpy array of bounding boxes \n        return : a string for each object in new line: <object-class> <x> <y> <width> <height>\n        \"\"\"\n        txt=''\n        for index,l in enumerate(bboxes):\n            l = [str(x) for x in l[:4]]\n            l = ' '.join(l)\n            txt +=  '0' +' ' + l + '\\n'\n        return txt\n\n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self,index):\n        row = self.df.iloc[index]\n        path = row['image_path']\n        img = cv2.imread(path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        aug_index = row['aug_index']\n        list_info = path.split('/')\n        image_name = list_info[-2] + '_' + list_info[-1].split('.')[0]\n        box = row['bboxes']\n        bounding_box = self.coord_to_box(box, img)\n\n        if self.transform is not None and self.mode == 'train':\n            res = self.transform(image=img, bboxes=bounding_box)\n            img = res['image']\n            bounding_box = res['bboxes']\n            \n        box_yolo_format = self.bbox_to_txt(bounding_box)\n        return img, box_yolo_format, image_name, aug_index\n","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.816624Z","iopub.execute_input":"2022-01-02T08:03:10.817131Z","iopub.status.idle":"2022-01-02T08:03:10.835227Z","shell.execute_reply.started":"2022-01-02T08:03:10.817088Z","shell.execute_reply":"2022-01-02T08:03:10.834283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_transforms(phase):\n    if not CFG.aug_box:\n        return None\n    if phase == 'train':\n        return A.Compose([\n                A.HorizontalFlip(p=0.3),\n                A.OneOf([\n#                     A.FDA(reference_images=FDA_image,p=0.5),\n                    HE_HSV(0.75),\n                    CLAHE_HSV(0.75),\n                    Gamma_enhancement(0.75)\n                ], p=0.75),\n#                 A.ShiftScaleRotate(scale_limit = 0, rotate_limit=30, p=0.3, border_mode=0)\n            ], bbox_params=A.BboxParams(format='yolo' , min_visibility=0.4,min_area=500))\n    else:\n        return None\n#         return A.Compose([])","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.836302Z","iopub.execute_input":"2022-01-02T08:03:10.836545Z","iopub.status.idle":"2022-01-02T08:03:10.851373Z","shell.execute_reply.started":"2022-01-02T08:03:10.836518Z","shell.execute_reply":"2022-01-02T08:03:10.850520Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = {\n    phase: AUG_DATASET(eval(phase), mode=phase, transform = get_transforms(phase=phase)) for phase in ['train','valid']\n}\n\ndataloader = {\n    phase: Data.DataLoader(dataset=dataset[phase], num_workers=CFG.worker, batch_size=CFG.batch_size, shuffle=False, drop_last=False,\\\n                           pin_memory = False) for phase in ['train','valid']\n}","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.852632Z","iopub.execute_input":"2022-01-02T08:03:10.854256Z","iopub.status.idle":"2022-01-02T08:03:10.869928Z","shell.execute_reply.started":"2022-01-02T08:03:10.854210Z","shell.execute_reply":"2022-01-02T08:03:10.868677Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for phase in ['train','valid']:\n    for aug_img, aug_box, image_name, aug_index in progress_bar(dataloader[phase]):\n        for idx, image in enumerate(aug_img):\n            name = image_name[idx]\n            aug_index_name = aug_index[idx]\n            new_name =  \"{}_{}\".format(name,aug_index_name)\n            image = aug_img[idx]\n            box = aug_box[idx]\n\n            path_txt = CFG.path_dataset + \"/\" + \"labels\" + \"/\" + phase + \"/\" + new_name + \".txt\"\n            path_jpg = CFG.path_dataset + \"/\" + \"images\" + \"/\" + phase + \"/\" + new_name + \".jpg\"\n            is_path = os.path.exists(path_jpg)\n            image = image.numpy()\n            cv2.imwrite(path_jpg, image[...,::-1])\n            txt_file = open(path_txt, \"w\")\n            txt_file.write(box)\n            txt_file.close()\n        break\n    break","metadata":{"execution":{"iopub.status.busy":"2022-01-02T08:03:10.871465Z","iopub.execute_input":"2022-01-02T08:03:10.872426Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}