{"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":"code","source":"import cv2\nimport os\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import DataLoader, Dataset\n\nimport albumentations\nfrom albumentations import pytorch as AT\n\nIMAGE_SIZE = 512\n\nrgby_mean = [0.08123, 0.05293, 0.05398, 0.08153]\nrgby_std  = [0.13028, 0.08611, 0.14256, 0.12620]\n\n\ntrain_transform = albumentations.Compose([\n    albumentations.ToFloat(max_value=65535.0),\n    albumentations.Resize(IMAGE_SIZE, IMAGE_SIZE),\n    albumentations.RandomRotate90(p=0.5),\n    albumentations.Transpose(p=0.5),\n    albumentations.Flip(p=0.5),    \n    albumentations.OneOf([\n        albumentations.ElasticTransform(alpha=1, sigma=20, alpha_affine=10),\n        albumentations.GridDistortion(num_steps=6, distort_limit=0.1),\n        albumentations.OpticalDistortion(distort_limit=0.05, shift_limit=0.05),\n    ], p=0.2), \n    albumentations.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05, p=0.5),\n    albumentations.core.composition.PerChannel(\n        albumentations.OneOf([\n            albumentations.MotionBlur(p=.05),\n            albumentations.MedianBlur(blur_limit=3, p=.05),\n            albumentations.Blur(blur_limit=3, p=.05),])\n        , p=1.0),\n    albumentations.OneOf([\n        albumentations.CoarseDropout(max_holes=16, max_height=IMAGE_SIZE//16, max_width=IMAGE_SIZE//16, fill_value=0, p=0.5),\n        albumentations.GridDropout(ratio=0.09, p=0.5),\n        albumentations.Cutout(num_holes=8, max_h_size=IMAGE_SIZE//16, max_w_size=IMAGE_SIZE//16, p=0.2),\n    ], p=0.5), \n    albumentations.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.5),\n    AT.ToTensorV2(),\n    ],\n    additional_targets={\n        'r': 'image',\n        'g': 'image',\n        'b': 'image',\n        'y': 'image',\n    }\n    )\n    \n    \ntest_transform = albumentations.Compose([\n    albumentations.ToFloat(max_value=65535.0),\n    albumentations.Resize(IMAGE_SIZE, IMAGE_SIZE),\n    AT.ToTensorV2(),\n    ],\n    additional_targets={\n        'r': 'image',\n        'g': 'image',\n        'b': 'image',\n        'y': 'image',\n    }\n    )\n    \n        \ntta_transform = albumentations.Compose([\n    albumentations.ToFloat(max_value=65535.0),\n    albumentations.RandomRotate90(p=0.5),\n    albumentations.Transpose(p=0.5),\n    albumentations.Flip(p=0.5),\n    albumentations.Resize(IMAGE_SIZE, IMAGE_SIZE),\n    AT.ToTensorV2(),\n    ],\n    additional_targets={\n        'r': 'image',\n        'g': 'image',\n        'b': 'image',\n        'y': 'image',\n    }\n    )\n    \n\nclass ImageDataset(Dataset):\n    \n    def __init__(self, df, data_path='../input', transform = train_transform): \n        self.df = df \n        self.data_path = data_path\n        self.imgs_path = self.df['folder']\n        self.imgs_name = self.df['ID']\n        self.labels = self.df['Label']\n        self.transform = transform\n\n\n    def __len__(self):\n        return self.df.shape[0]\n\n    def __getitem__(self, idx):\n        img_path = os.path.join(self.data_path, self.imgs_path[idx], self.imgs_name[idx])\n        img_r = cv2.imread(img_path + '_red.png', -1)\n        img_g = cv2.imread(img_path + '_green.png', -1)\n        img_b = cv2.imread(img_path + '_blue.png', -1)\n        img_rgb = np.dstack((img_r, img_g, img_b))\n        img_y = cv2.imread(img_path + '_yellow.png', 3)\n        # transform\n        trans_img = self.transform(image=img_rgb, y=img_y)\n        img_rgb, img_y = trans_img['image'], trans_img['y']\n        img_rgby = np.concatenate([img_rgb, img_y], axis=0)[:4]\n        # label\n        onehot = np.zeros(18) # 0-17 18:neg\n        label = self.labels[idx]\n        if label != '18':\n            label = list(set(label.split('|')))\n            for i in label:\n                onehot[int(i)] += 1\n        label = onehot\n        return img_rgby, label\n        \n        \nclass HPAMixup(Dataset):\n    def __init__(self, dataset, num_class, max_mix=3, prob=0.1):\n        self.dataset = dataset\n        self.num_class = num_class\n        self.max_mix = int(max_mix)\n        self.prob = prob\n        self.data_size = len(self)\n\n    def __getitem__(self, idx):\n        img, label = self.dataset[idx]\n        if self.max_mix <= 1:\n            return img, label\n        \n        for i in range(self.max_mix-1):\n            if torch.rand(1)[0] < self.prob:\n                rand_idx = torch.randint(self.data_size,(1,))[0].numpy()\n                img_aug, label_aug = self.dataset[rand_idx]\n                img += img_aug\n                label = label + label_aug \n                \n        label = label > 0 # binary label\n        label = label.astype(np.float)\n\n        return img, label\n\n    def __len__(self):\n        return len(self.dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# usage\n# \"\"\"\n# trainset = ImageDataset(train_df.iloc[tr_folds[fold]].reset_index(), base_dir, train_transform)\n# trainset_mixup = HPAMixup(trainset, CLASSES_NUM, max_mix=3, prob=0.1) \n# train_loader = torch.utils.data.DataLoader(trainset_mixup, batch_size=BATCH_SIZE, num_workers=16, shuffle=True, drop_last=True, worker_init_fn=worker_init_fn)\n# \"\"\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}