{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport torch\nfrom torchvision import models\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!git clone https://github.com/cocodataset/cocoapi.git\n%cd /kaggle/working/cocoapi/PythonAPI\n!python setup.py build_ext install","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# %%shell\n%cd /kaggle/working/\n\n!pip install cython\n# Install pycocotools, the version by default in Colab\n# has a bug fixed in https://github.com/cocodataset/cocoapi/pull/354\n!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'\n# Download TorchVision repo to use some files from\n# references/detection\n!git clone https://github.com/pytorch/vision.git\n%cd vision\n!git checkout v0.3.0\n\n!cp references/detection/utils.py ../\n!cp references/detection/transforms.py ../\n!cp references/detection/coco_eval.py ../\n!cp references/detection/engine.py ../\n!cp references/detection/coco_utils.py ../","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from engine import train_one_epoch, evaluate\nimport utils\nimport transforms as T","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('/kaggle/input/global-wheat-detection/train.csv')\nprint(train_df.shape)\n# test_df = pd.read_csv('/kaggle/input/global-wheat-detection/test.csv')\n# print(test_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_ids = train_df['image_id'].unique()\nvalid_ids = img_ids[-300:]\ntrain_ids = img_ids[:-300]\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\nvalid_df.shape,train_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## A TEST TO IMROVE THE SCORE - NORMALIZATION OF THE IMAGE PIXEL VALUES (DIVISION BY 255)\n\n# img = cv2.imread(\"/kaggle/input/global-wheat-detection/test/2fd875eaa.jpg\")\n# img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n# plt.figure()\n# plt.imshow(img)\n# i = img/255\n# plt.figure()\n# plt.imshow(i)\n# print(img,i)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GWDataset(torch.utils.data.Dataset):\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.img_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index):\n        img_idx = self.img_ids[index]\n        img_name = str(img_idx+'.jpg')\n        # load images ad masks\n        img_path = os.path.join(self.image_dir, img_name)\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).astype(np.float32)\n        img = img/255.0\n       # get bounding box coordinates for each mask\n        num_bbxs = len(self.df[self.df['image_id']==img_idx])\n        bbxs = self.df[self.df['image_id']==img_idx]\n        boxes = []\n        area = []\n#         print(bbxs)\n        for t in range(num_bbxs):\n            l = bbxs.iloc[t]['bbox'].split(',')\n#             print(l)\n            xmin,ymin,w,h = float(l[0][1:]),float(l[1][1:]),float(l[2][1:]),float(l[3][1:-1])\n            xmax = xmin+w\n            ymax = ymin+h\n            area.append(w*h)\n            boxes.append([xmin, ymin, xmax, ymax])\n\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        # there is only one class\n        labels = torch.ones((num_bbxs,), dtype=torch.int64)\n\n        imag_id = torch.tensor([index])\n        # suppose all instances are not crowd\n        area = torch.as_tensor(area, dtype=torch.float32)\n        iscrowd = torch.zeros((num_bbxs,), dtype=torch.int64)\n\n        target = {}\n        target[\"boxes\"] = boxes\n        target[\"labels\"] = labels\n        target[\"image_id\"] = imag_id\n        target[\"area\"] = area\n        target[\"iscrowd\"] = iscrowd\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n\n    def __len__(self):\n        return (self.img_ids.shape[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transform(train):\n    transforms = []\n    # converts the image, a PIL image, into a PyTorch Tensor\n    transforms.append(T.ToTensor())\n    if train:\n        # during training, randomly flip the training images\n        # and ground-truth for data augmentation\n        transforms.append(T.RandomHorizontalFlip(0.5))\n    return T.Compose(transforms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = GWDataset(valid_df,'/kaggle/input/global-wheat-detection/train/',get_transform(train = False))\ndataset[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"##  CHOICE BETWEEN MODEL PRETRAINED ON image_net vs Global Wheat Detection Challenge\n# model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\nmodel = torch.load(\"/kaggle/input/gwd-model/fasterrcnn_resnet50_fpn.pth\",map_location='cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = 2  # 1 class (wheat) + background\n\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get_train_transform()=\ndef collate_fn(batch):\n    return tuple(zip(*batch))\nDIR_TRAIN = '/kaggle/input/global-wheat-detection/train/'\ntrain_dataset = GWDataset(train_df, DIR_TRAIN,get_transform(train = True))\nvalid_dataset = GWDataset(valid_df, DIR_TRAIN, get_transform(train = False))\n\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nmodel.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.0005, momentum=0.9, weight_decay=0.0005)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.1)\nnum_epochs = 2\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs = 4\n\nfor epoch in range(num_epochs):\n    # train for one epoch, printing every 10 iterations\n    train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=10)\n    # update the learning rate\n    lr_scheduler.step()\n    # evaluate on the test dataset\n    evaluate(model, valid_data_loader, device=device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# pick one image from the test set\nimg, _ = valid_dataset[0]\n# put the model in evaluation mode\nmodel.eval()\nwith torch.no_grad():\n    prediction = model([img.to(device)])\nsample = valid_dataset[0][0].permute(1,2,0).numpy()\nboxes = prediction[0]['boxes'].cpu().numpy().astype(np.int32)\n# boxe = boxes.reshape((4,-1))\nscores = prediction[0]['scores'].cpu().numpy()\nboxes.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"scores","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# plt.imshow(sample)\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\ncolor = (220,0,0)\nfor i in range(len(boxes)):\n#     print(type(box[0]))\n    if scores[i]>0.97:\n        cv2.rectangle(img,(int(boxes[i][0]), int(boxes[i][1])),(int(boxes[i][2]), int(boxes[i][3])),color, 5)\nax.set_axis_off()\nax.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model, '/kaggle/working/fasterrcnn_resnet50_fpn_new.pth')\ntorch.save(model.state_dict(), '/kaggle/working/fasterrcnn_resnet50_fpn_statedict.pth')","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}