{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# 1.4 data augmentation\nimport imgaug as ia\nfrom imgaug import augmenters as iaa\n\n\ndef augmentation_generator(yolo_dataset):\n    for batch in yolo_dataset:\n        # conversion tensor->numpy\n        img = batch[0].numpy()\n        boxes = batch[1].numpy()\n        # conversion bbox numpy->ia object\n        ia_boxes = []\n        for i in range(img.shape[0]):\n            ia_bbs = [ia.BoundingBox(x1=bb[0],\n                                     y1=bb[1],\n                                     x2=bb[2],\n                                     y2=bb[3]) for bb in boxes[i]\n                      if (bb[0] + bb[1] + bb[2] + bb[3] > 0)]\n            ia_boxes.append(ia.BoundingBoxesOnImage(ia_bbs, shape=(512, 512)))\n        # data augmentation\n        seq = iaa.Sequential([\n            iaa.Fliplr(0.5),\n            iaa.Flipud(0.5),\n            iaa.Multiply((0.4, 1.6)),  # change brightness\n            # iaa.ContrastNormalization((0.5, 1.5)),\n            # iaa.Affine(translate_px={\"x\": (-100,100), \"y\": (-100,100)}, scale=(0.7, 1.30))\n        ])\n        # seq = iaa.Sequential([])\n        seq_det = seq.to_deterministic()\n        img_aug = seq_det.augment_images(img)\n        img_aug = np.clip(img_aug, 0, 1)\n        boxes_aug = seq_det.augment_bounding_boxes(ia_boxes)\n        # conversion ia object -> bbox numpy\n        for i in range(img.shape[0]):\n            boxes_aug[i] = boxes_aug[i].remove_out_of_image().clip_out_of_image()\n            for j, bb in enumerate(boxes_aug[i].bounding_boxes):\n                boxes[i, j, 0] = bb.x1\n                boxes[i, j, 1] = bb.y1\n                boxes[i, j, 2] = bb.x2\n                boxes[i, j, 3] = bb.y2\n        # conversion numpy->tensor\n        batch = (tf.convert_to_tensor(img_aug), tf.convert_to_tensor(boxes))\n        # batch = (img_aug, boxes)\n        yield batch\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"aug_train_db = augmentation_generator(train_db)\ndb_visualize(aug_train_db)","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}