{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\nimport os\nimport cv2\nimport timm\nimport torch\nfrom torch import nn\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset,DataLoader\n\nfrom albumentations import (\n    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize\n)\n\nfrom albumentations.pytorch import ToTensorV2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG = {\n    'model_arch': 'tf_efficientnet_b5_ns',\n    'img_size': 456,\n    'valid_bs': 32,\n    'num_workers': 4,\n    'device': 'cuda:0',\n    'tta': 4,\n    'used_epochs': [7,9,9,9,9],\n    'weights': [1,1,1,1,1]\n}\nsubmission = pd.read_csv('../input/cassava-leaf-disease-classification/sample_submission.csv')\nsubmission.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_img(path):\n    '''使用 opencv 加载图片.\n    由于历史原因，opencv 读取的图片格式是 bgr\n    Args:\n        path : str  图片文件路径 e.g '../data/train_img/1.jpg'\n    '''\n    img_bgr = cv2.imread(path)\n    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n    return img_rgb\nimg = get_img('../input/cassava-leaf-disease-classification/train_images/1000015157.jpg')\nplt.imshow(img)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    '''木薯叶比赛数据加载类\n    Attributes:\n        __len__ : 数据的样本个数.\n        __getitem__ : 索引函数.\n    '''\n    def __init__(\n            self,\n            df,\n            data_root,\n            transforms=None,\n            output_label=True,\n            one_hot_label=False,\n            do_fmix=False,\n            fmix_params={\n                'alpha': 1.,\n                'decay_power': 3.,\n                'shape': (512, 512),\n                'max_soft': 0.3,\n                'reformulate': False\n            },\n            do_cutmix=False,\n            cutmix_params={\n                'alpha': 1,\n            }):\n        '''\n        Args:\n            df : DataFrame , 样本图片的文件名和标签\n            data_root : str , 图片所在的文件路径，绝对路径\n            transforms : object , 图片增强\n            output_label : bool , 是否输出标签\n            one_hot_label : bool , 是否进行 onehot 编码\n            do_fmix : bool , 是否使用 fmix\n            fmix_params :dict , fmix 的参数 {'alpha':1.,'decay_power':3.,'shape':(256,256),'max_soft':0.3,'reformulate':False}\n            do_cutmix : bool, 是否使用 cutmix\n            cutmix_params : dict , cutmix 的参数 {'alpha':1.}\n        Raises:\n\n        '''\n        super().__init__()\n        self.df = df.reset_index(drop=True).copy()  # 重新生成索引\n        self.transforms = transforms\n        self.data_root = data_root\n        self.do_fmix = do_fmix\n        self.fmix_params = fmix_params\n        self.do_cutmix = do_cutmix\n        self.cutmix_params = cutmix_params\n        self.output_label = output_label\n        self.one_hot_label = one_hot_label\n        if output_label:\n            self.labels = self.df['label'].values\n            if one_hot_label:\n                self.labels = np.eye(self.df['label'].max() +\n                                     1)[self.labels]  # 使用单位矩阵生成 onehot 编码\n\n    def __len__(self):\n        return self.df.shape[0]\n\n    def __getitem__(self, index):\n        '''\n        Args:\n            index : int , 索引\n        Returns:\n            img, target(optional)\n        '''\n        if self.output_label:\n            target = self.labels[index]\n\n        img = get_img(\n            os.path.join(self.data_root,\n                         self.df.loc[index]['image_id']))  # 拼接地址，加载图片\n\n        if self.transforms:  # 使用图片增强\n            img = self.transforms(image=img)['image']\n\n        if self.do_fmix and np.random.uniform(\n                0., 1., size=1)[0] > 0.5:  # 50% 概率触发 fmix 数据增强\n\n            with torch.no_grad():\n                lam, mask = sample_mask(\n                    **self.fmix_params)  # 可以考虑魔改，使用 clip 规定上下限制\n\n                fmix_ix = np.random.choice(self.df.index,\n                                           size=1)[0]  # 随机选择待 mix 的图片\n                fmix_img = get_img(\n                    os.path.join(self.data_root,\n                                 self.df.loc[fmix_ix]['image_id']))\n\n                if self.transforms:\n                    fmix_img = self.transforms(image=fmix_img)['image']\n\n                mask_torch = torch.from_numpy(mask)\n\n                img = mask_torch * img + (1. - mask_torch) * fmix_img  # mix 图片\n\n                rate = mask.sum() / float(img.size)  # 获取 mix 的 rate\n                target = rate * target + (\n                    1. - rate) * self.labels[fmix_ix]  # target 进行 mix\n\n        if self.do_cutmix and np.random.uniform(\n                0., 1., size=1)[0] > 0.5:  # 50% 概率触发 cutmix 数据增强\n            with torch.no_grad():\n                cmix_ix = np.random.choice(self.df.index, size=1)[0]\n                cmix_img = get_img(\n                    os.path.join(self.data_root,\n                                 self.df.loc[cmix_ix]['image_id']))\n                if self.transforms:\n                    cmix_img = self.transforms(image=cmix_img)['image']\n\n                lam = np.clip(\n                    np.random.beta(self.cutmix_params['alpha'],\n                                   self.cutmix_params['alpha']), 0.3, 0.4)\n                bbx1, bby1, bbx2, bby2 = rand_bbox(cmix_img.shape[:2], lam)\n\n                img[:, bbx1:bbx2, bby1:bby2] = cmix_img[:, bbx1:bbx2,\n                                                        bby1:bby2]\n\n                rate = 1 - ((bbx2 - bbx1) *\n                            (bby2 - bby1) / float(img.size))  # 获取 mix 的 rate\n                target = rate * target + (\n                    1. - rate) * self.labels[cmix_ix]  # target 进行 mix\n\n        if self.output_label:\n            return img, target\n        else:\n            return img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_inference_transforms():\n    '''测试阶段的数据增强TTA使用\n    '''\n    return Compose([\n            RandomResizedCrop(CFG['img_size'], CFG['img_size']),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            #HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n            #RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassvaImgClassifier(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        n_features = self.model.classifier.in_features\n        self.model.classifier = nn.Linear(n_features, n_class)\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def inference_one_epoch(model, data_loader, device):\n    model.eval()\n    image_preds_all = []\n    pbar = tqdm(enumerate(data_loader), total=len(data_loader))\n    print(pbar)\n    for step, (imgs) in pbar:\n        imgs = imgs.to(device).float()\n        image_preds = model(imgs) \n        image_preds_all += [torch.softmax(image_preds, 1).detach().cpu().numpy()]\n    image_preds_all = np.concatenate(image_preds_all, axis=0)\n    return image_preds_all","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = pd.DataFrame()\ntest['image_id'] = list(os.listdir('../input/cassava-leaf-disease-classification/test_images/'))\ntest_ds = CassavaDataset(test, '../input/cassava-leaf-disease-classification/test_images/', transforms=get_inference_transforms(), output_label=False)\n\ntst_loader = torch.utils.data.DataLoader(\n            test_ds, \n            batch_size=CFG['valid_bs'],\n            num_workers=CFG['num_workers'],\n            shuffle=False,\n            pin_memory=False,\n        )\ndevice = torch.device(CFG['device'])\nmodel = CassvaImgClassifier(CFG['model_arch'], 5).to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_preds = []\ntst_preds = []\n\n# 每个 fold 的模型选择最好的 epoch 的结果，做 3 次 tta 融合得到最终结果。\nfor fold, epoch in enumerate(CFG['used_epochs']):    \n    model.load_state_dict(torch.load('../input/new60-b5/{}_fold_{}_{}'.format(CFG['model_arch'], fold, epoch)))\n\n    with torch.no_grad():\n        for _ in range(CFG['tta']):\n            tst_preds += [CFG['weights'][fold]/sum(CFG['weights'])/CFG['tta']*inference_one_epoch(model, tst_loader, device)]\n\ntst_preds = np.sum(tst_preds, axis=0) \n\ndel model\ntorch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test['label'] = np.argmax(tst_preds, axis=1)\ntest.to_csv('submission.csv', index=False)\ntest.head()","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}