{"cells":[{"metadata":{},"cell_type":"markdown","source":"This idea copied from [this kernel](https://www.kaggle.com/abhinand05/vision-transformer-vit-tutorial-baseline)"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ../input/pytorchlightning/tensorboard-2.2.0-py3-none-any.whl\n!pip install ../input/pytorchlightning/pytorch_lightning-0.9.0-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport sys\nsys.path.append(\"../input/timm-pytorch-models\")\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedKFold\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.metrics.functional.classification import accuracy\n\nimport timm\n#timm.list_models(\"vit*\")\n\nimport torch\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch import nn\n\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensor","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Input Constants"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_PATH = \"../input/vit-base-models-pretrained-pytorch/jx_vit_base_p32_384-830016f5.pth\"\nCLASSES = 5\nIMG_DIR = \"../input/cassava-leaf-disease-classification/train_images\"\nIMAGE_SIZE = 384 \nTRAIN_FILE = \"../input/cassava-leaf-disease-classification/train.csv\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavViT(pl.LightningModule):\n  def __init__(self):\n    super().__init__()\n    self.loss_fn = nn.CrossEntropyLoss()\n    self.model = timm.create_model(\"vit_base_patch32_384\",pretrained=False)\n    self.model.load_state_dict(torch.load(MODEL_PATH))\n    self.model.head.out_features = CLASSES\n  \n  def forward(self,x):\n    return self.model(x)\n  \n  def configure_optimizers(self):\n    LR = 1e-5\n    return torch.optim.Adam(self.parameters(),lr=LR)\n\n  def training_step(self,batch,batch_idx):\n    x,y = batch[\"x\"],batch[\"y\"]\n    y_hat = self(x)\n    loss = self.loss_fn(y_hat,y)\n    self.log(\"train_loss\",loss)\n    acc = accuracy(y_hat,y)\n    self.log(\"train_acc\",acc,on_epoch=True,prog_bar=True)\n    return loss \n\n  def validation_step(self,batch,batch_idx):\n    x,y = batch[\"x\"],batch[\"y\"]\n    y_hat = self(x)\n    loss = self.loss_fn(y_hat,y)\n    acc = accuracy(y_hat,y)\n    self.log(\"val_acc\",acc,on_epoch=True,prog_bar=True)\n    self.log(\"val_loss\",loss,prog_bar=True) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset Generator"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class CassavData(Dataset):\n  def __init__(self,path,image_ids,labels,transform):\n    super().__init__()\n    self.path = path\n    self.image_ids = image_ids\n    self.labels = labels\n    self.transform = transform\n\n  def __len__(self):\n    return len(self.image_ids)\n  \n  def __getitem__(self,item):\n    img_id = str(self.image_ids[item])\n    label = self.labels[item]\n    img = cv2.imread(os.path.join(self.path,img_id))\n    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n    img = self.transform(image=img)[\"image\"]\n    return {\n        \"x\":img,\n        \"y\":label\n    }\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Module"},{"metadata":{"trusted":true},"cell_type":"code","source":"  \nclass CassavDataModule(pl.LightningDataModule):\n  def __init__(self):\n    super().__init__()\n    self.train_transform = A.Compose([\n                                 A.Resize(IMAGE_SIZE,IMAGE_SIZE),\n                                 A.HorizontalFlip(),\n                                 A.VerticalFlip(),\n                                 A.RandomCrop(IMAGE_SIZE,\n                                                   IMAGE_SIZE),\n                                 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n                                 ToTensor()])\n    self.test_transform = A.Compose([\n                                 A.Resize(IMAGE_SIZE,IMAGE_SIZE),\n                                 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n                                 ToTensor()])\n  \n  def setup(self,stage=None):\n    dfx = pd.read_csv(TRAIN_FILE)\n    dfx[\"kfold\"] = -1\n    dfx = dfx.sample(frac=1).reset_index(drop=True)\n    skfold = StratifiedKFold(n_splits=5)\n    for i,(t_idx,v_idx) in enumerate(skfold.split(dfx.image_id,y=dfx.label)):\n      dfx.loc[v_idx,\"kfold\"]=i\n    validation = dfx.loc[dfx.kfold==1]\n    train = dfx.loc[dfx.kfold!=1]\n    self.train_dataset = CassavData(path=IMG_DIR,\n                               image_ids = train.image_id.values,\n                               labels = train.label.values,\n                               transform = self.train_transform)\n    self.val_dataset = CassavData(path=IMG_DIR,\n                             image_ids = validation.image_id.values,\n                               labels = validation.label.values,\n                               transform = self.test_transform)     \n  \n  def train_dataloader(self):\n    return DataLoader(dataset=self.train_dataset,\n                      batch_size=16,\n                      num_workers=4,\n                      drop_last=True,\n                      shuffle=True)\n  \n  def val_dataloader(self):\n    return DataLoader(dataset=self.val_dataset,\n                      batch_size=16,\n                      num_workers=4,\n                      drop_last=True,\n                      shuffle=False) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training & Save model"},{"metadata":{"trusted":true},"cell_type":"code","source":"#dm = CassavDataModule()\n#net = CassavViT() \n#trainer = pl.Trainer(gpus=-1)\n#trainer.fit(model = net,\n#            datamodule = dm) \n#trainer.save_checkpoints(\"ViT_model.pth\",\n#                         max_epochs=20,)\nclass CassavTestData(Dataset):\n  def __init__(self,path,image_ids,transform):\n    super().__init__()\n    self.path = path\n    self.image_ids = image_ids\n\n    self.transform = transform\n\n  def __len__(self):\n    return len(self.image_ids)\n  \n  def __getitem__(self,item):\n    img_id = str(self.image_ids[item])\n    img = cv2.imread(os.path.join(self.path,img_id))\n    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n    img = self.transform(image=img)[\"image\"]\n    return {\n        \"x\":img,\n    }\n\n\nTEST_IMAGE_DIRS = \"../input/cassava-leaf-disease-classification/test_images/\"\ntest = pd.read_csv(\"../input/cassava-leaf-disease-classification/sample_submission.csv\")\ntest_transform = A.Compose([\n                                 A.Resize(IMAGE_SIZE,IMAGE_SIZE),\n                                 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n                                 ToTensor()])\ntest_dataset = CassavTestData(path = TEST_IMAGE_DIRS,\n                              image_ids = test.image_id.values,transform = test_transform)\ntest_loader = DataLoader(test_dataset,\n                        batch_size=32,\n                        )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#loading the best checkpoints to model\nbest_checkpoints = \"../input/vit-cassava-trained/ViT_model.pth\"\npretrained_model = CassavViT.load_from_checkpoint(checkpoint_path = best_checkpoints)\npretrained_model = pretrained_model.to(\"cuda\")\npretrained_model.eval()\npretrained_model.freeze()\n\nfin_out = []\nfor data in test_loader:\n    y_hat = pretrained_model(data[\"x\"].to(\"cuda\"))\n    y_hat = torch.argmax(y_hat,dim=1)\n    fin_out.extend(y_hat.cpu().detach().numpy().tolist())\ntest[\"label\"] = fin_out\ntest[[\"image_id\",\"label\"]].to_csv(\"submission.csv\",index=False)\ntest.head()\n","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}