{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"Hi, everyone！\nThis kernel is based on the work [WBF over TTA](https://www.kaggle.com/shonenkov/wbf-over-tta-single-model-efficientdet) of [Alex shonenkov](https://www.kaggle.com/shonenkov).Thanks for his sharing!\n\nI just add another transform and it gives me a slight improvement.\n\nHope you enjoy it！","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import albumentations as A\n\nclass BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 512\n    \n    def augment(self,image):\n        raise NotImplementedError\n        \n    def batch_augment(self,images):\n        raise NotImplementedError\n        \n    def deaugment_boxes(self,boxes):\n        raise NotImplementedError\n        \nclass TTAVerticalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self,image):\n        return image.flip(1)\n    \n    def batch_augment(self,images):\n        return images.flip(2)\n    \n    def deaugment_boxes(self,boxes):\n        boxes[:,[1,3]] = self.image_size - boxes[:,[3,1]]\n        return boxes\n    \nclass TTAHorizontalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return image.flip(2)\n    \n    def batch_augment(self, images):\n        return images.flip(3)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]\n        return boxes\n    \nclass TTARotate90(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 1, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 1, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = self.image_size - boxes[:, [1,3]]\n        res_boxes[:, [1,3]] = boxes[:, [2,0]]\n        return res_boxes\n    \nclass TTAHSV_or_RBC(BaseWheatTTA):\n    \"\"\"HSV or RBC\"\"\"\n    transform = A.OneOf([\n                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit= 0.2, \n                                     val_shift_limit=0.2,p=0.5),\n                A.RandomBrightnessContrast(brightness_limit=0.2, \n                                           contrast_limit=0.2, p=0.5),\n            ],p=1)\n    \n    def augment(self,image):\n        image = image.permute(1,2,0).cpu().numpy()\n        image = self.transform(image=image)['image']\n        return ToTensorV2()(image=image)['image']\n\n    def batch_augment(self, images):\n        batch_size = len(images)\n        for i in range(batch_size):\n            images[i] = self.augment(images[i])\n        return images\n        \n    def deaugment_boxes(self,boxes):\n        return boxes\n    \nclass TTACompose(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def augment(self, image):\n        for transform in self.transforms:\n            image = transform.augment(image)\n        return image\n    \n    def batch_augment(self, images):\n        for transform in self.transforms:\n            images = transform.batch_augment(images)\n        return images\n    \n    def prepare_boxes(self, boxes):\n        result_boxes = boxes.copy()\n        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)\n        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)\n        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)\n        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)\n        return result_boxes\n    \n    def deaugment_boxes(self, boxes):\n        for transform in self.transforms[::-1]:\n            boxes = transform.deaugment_boxes(boxes)\n        return self.prepare_boxes(boxes)","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}