{"cells":[{"metadata":{},"cell_type":"markdown","source":"This is notebook is based on [CassavaLeaf ViT baseline Notebook](https://www.kaggle.com/szuzhangzhi/vision-transformer-vit-cuda-as-usual) and Adjusted for Plant pathology 2021 dataset."},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\n\npackage_path = '../input/vision-transformer-pytorch/VisionTransformer-Pytorch'\nsys.path.append(package_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\n\nimport time\nimport datetime\nimport copy\nimport matplotlib.pyplot as plt\nimport json\nimport seaborn as sns\nimport cv2\nimport albumentations as albu\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold, train_test_split\n\n\n# ALBUMENTATIONS\nimport albumentations as albu\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\n\n# ADAMP\n# from adamp import AdamP","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_DIR=\"../input/plant-pathology-2021-fgvc8/\"\nTRAIN_IMAGES_DIR = os.path.join(BASE_DIR,'train_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(BASE_DIR, 'train.csv'))\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Count of training images {0}\".format(len(os.listdir(TRAIN_IMAGES_DIR))))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{BASE_DIR}/train.csv')\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class_name = train_df['labels'].value_counts().index\nclass_count = train_df['labels'].value_counts().values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.labels.value_counts()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Counting target values.\n\ntarg_cts=train_df.labels.value_counts()\nfig = plt.figure(figsize=(12,6))\nsns.barplot(y=targ_cts.sort_values(ascending=False).index, x=targ_cts.sort_values(ascending=False).values, palette='summer')\nplt.title('Target Distribution')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.preprocessing import LabelEncoder\nle = LabelEncoder()\n\nle.fit(train_df.labels)\ntrain_df['labels'] = le.transform(train_df.labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def visualize_images(image_ids, labels):\n    plt.figure(figsize=(16, 12))\n    \n    for idx, (image_id, label) in enumerate(zip(image_ids, labels)):\n        plt.subplot(3, 3, idx+1)\n        \n        image = cv2.imread(os.path.join(TRAIN_IMAGES_DIR, image_id))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        plt.imshow(image)\n        plt.title(f\"Class: {label}\", fontsize=12)\n        plt.axis(\"off\")\n        \n    plt.show()\n    \n\ndef plot_augmentation(image_id, transform):\n    plt.figure(figsize=(16, 4))\n    \n    img = cv2.imread(os.path.join(TRAIN_IAMGES_DIR, image_id))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    \n    plt.subplot(1, 3, 1)\n    plt.imshow(img)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 2)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 3)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    \ndef visualize(images, transform):\n    '''\n    Plot images and their transformations\n    '''\n    fig = plt.figure(figsize=(32, 16))\n    \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+1, xticks=[], yticks=[])\n        plt.imshow(im)\n        \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+6, xticks=[], yticks=[])\n        plt.imshow(transform(image=im)['image'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def visualize_images(image_ids, labels):\n    plt.figure(figsize=(16, 12))\n    \n    for idx, (image_id, label) in enumerate(zip(image_ids, labels)):\n        plt.subplot(3, 3, idx+1)\n        \n        image = cv2.imread(os.path.join(TRAIN_IMAGES_DIR, image_id))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        plt.imshow(image)\n        plt.title(f\"Class: {label}\", fontsize=12)\n        \n        plt.axis(\"off\")\n        \n    plt.show()\n    \n\ndef plot_augmentation(image_id, transform):\n    plt.figure(figsize=(16, 4))\n    \n    img = cv2.imread(os.path.join(TRAIN_IAMGES_DIR, image_id))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    \n    plt.subplot(1, 3, 1)\n    plt.imshow(img)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 2)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 3)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    \ndef visualize(images, transform):\n    '''\n    Plot images and their transformations\n    '''\n    fig = plt.figure(figsize=(32, 16))\n    \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+1, xticks=[], yticks=[])\n        plt.imshow(im)\n        \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+6, xticks=[], yticks=[])\n        plt.imshow(transform(image=im)['image'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# CUSTOM DATASET CLASS\nclass PlantDataset(Dataset):\n    def __init__(\n        self, df:pd.DataFrame, imfolder:str, train:bool=True, transforms=None\n    ):\n        self.df = df\n        self.imfolder = imfolder\n        self.train = train\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        im_path = os.path.join(self.imfolder, self.df.iloc[index]['image'])\n        im = cv2.imread(im_path, cv2.IMREAD_COLOR)\n        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)\n        \n        if (self.transforms):\n            '''\n            When AlbumentationCompose, a dictionary with key 'image' is created\n            '''\n            im = self.transforms(image=im)['image']\n            \n        if (self.train):\n            label = self.df.iloc[index]['labels']\n            return im, label\n        else:\n            return im","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# AUGMENTATIONS\n\ntrain_augs = albu.Compose([\n    albu.RandomResizedCrop(height=384, width=384, p=1.0),\n    albu.HorizontalFlip(p=0.5),\n    albu.VerticalFlip(p=0.5),\n    albu.RandomBrightnessContrast(p=0.5),\n    albu.ShiftScaleRotate(p=0.5),\n    albu.Normalize(    \n        mean=[0.3, 0.3, 0.3],\n        std=[0.1, 0.1, 0.1],),\n    CoarseDropout(p=0.5),\n    Cutout(p=0.5),\n    ToTensorV2(),\n])\n\nvalid_augs = albu.Compose([\n    albu.Resize(height=384, width=384, p=1.0),\n    albu.Normalize(\n        mean=[0.3, 0.3, 0.3],\n        std=[0.1, 0.1, 0.1],),\n    ToTensorV2(),\n])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# DATA SPLIT\ntrain, valid = train_test_split(\n    train_df,\n    test_size=0.1,\n    random_state=42,\n    stratify=train_df.labels.values\n)\n\n# reset index on both dataframes\ntrain = train.reset_index(drop=True)\nvalid = valid.reset_index(drop=True)\n\n# targets in train,valid datasets\ntrain_targets = train.labels.values\nvalid_targets = valid.labels.values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# DEFINE PYTORCH CUSTOM DATASET\ntrain_dataset = PlantDataset(\n    df = train,\n    imfolder = TRAIN_IMAGES_DIR,\n    train = True,\n    transforms = train_augs\n)\n\nvalid_dataset = PlantDataset(\n    df = valid,\n    imfolder = TRAIN_IMAGES_DIR,\n    train=True,\n    transforms = valid_augs\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_image(img_dict):\n    image_tensor = img_dict[0]\n#     print(type(image_tensor))\n    target = img_dict[1]\n    print(target)\n    image = image_tensor.permute(1, 2, 0)\n    plt.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_image(train_dataset[7])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# MAKE PYTORCH DATALOADER\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size = 4,\n    num_workers = 4,\n    shuffle = True\n)\n\nvalid_loader = DataLoader(\n    valid_dataset,\n    batch_size = 4,\n    num_workers = 4,\n    shuffle = False\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# TRAIN\ndef train_model(datasets, dataloaders, model, criterion, optimizer, scheduler, num_epochs, device):\n    since = time.time()\n    \n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_acc = 0.0\n    \n    for epoch in range(num_epochs):\n        print('Epoch {}/{}'.format(epoch, num_epochs-1))\n        print('-' * 10)\n        \n        for phase in ['train', 'valid']:\n            if phase == 'train':\n                model.train()\n            else:\n                model.eval()\n                \n            running_loss = 0.0\n            running_corrects = 0.0\n            \n            for inputs, labels in dataloaders[phase]:\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n               \n                # Zero out the grads\n                optimizer.zero_grad()\n                \n                # Forward\n                # Track history in train mode\n                with torch.set_grad_enabled(phase == 'train'):\n                    model = model.to(device)\n                    outputs = model(inputs)\n                    _, preds = torch.max(outputs, 1) \n                    loss = criterion(outputs, labels)\n                    \n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n                \n                # Statistics\n                running_loss += loss.item()*inputs.size(0)\n                running_corrects += torch.sum(preds == labels.data)\n            \n            if phase == 'train':\n                scheduler.step()\n                \n            epoch_loss = running_loss / len(datasets[phase])\n            epoch_acc = running_corrects.double() / len(datasets[phase])\n            \n            print('{} Loss: {:.4f} Acc: {:.4f}'.format(\n                phase, epoch_loss, epoch_acc))\n            \n            if phase == 'valid' and epoch_acc > best_acc:\n                best_acc = epoch_acc\n                best_model_wts = copy.deepcopy(model.state_dict())\n        \n        print()\n    \n    time_elapsed = time.time() - since\n    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n    print('Best val Acc: {:.4f}'.format(best_acc))\n    \n    model.load_state_dict(best_model_wts)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from vision_transformer_pytorch import VisionTransformer\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)\n\ndatasets = {'train': train_dataset,\n            'valid': valid_dataset}\n\ndataloaders = {'train': train_loader,\n               'valid': valid_loader}\n\n# LOAD PRETRAINED ViT MODEL\nmodel = VisionTransformer.from_pretrained('ViT-B_16', num_classes=12)        \n\n# OPTIMIZER\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.001)\n# optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=0.001)\n# optimizer = AdamP(model.parameters(), lr=1e-4, weight_decay=0.001)\n\n# LEARNING RATE SCHEDULER\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)\n\ncriterion = nn.CrossEntropyLoss()\nnum_epochs = 2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# MODEL TRAIN\ntrained_model = train_model(datasets, dataloaders, model, criterion, optimizer, scheduler, num_epochs, device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Save the mode after training\ntorch.save(model.state_dict(), 'vit_b-16_2epoch.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}