{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Global Wheat Detection Faster R-CNN"},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom PIL import Image\nimport matplotlib.patches as patches\nimport matplotlib.pyplot as plt\nimport os\nimport ast","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import Dataset\nimport torchvision.transforms as T\nimport torch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class GlobalWheatDataset(Dataset):\n    def __init__(self, path, train_or_test_path, transforms):\n        self.path = path\n        self.train_or_test_path = train_or_test_path\n        self.transforms = transforms\n\n        self.df = pd.read_csv(path + 'train.csv')\n\n        self.ids = {v:k for k, v in enumerate(np.unique(self.df.image_id.values))}\n        self.imgs_list = list(sorted(os.listdir(os.path.join(path, train_or_test_path))))\n\n    def get_rectangles(self, idx):\n        id = self.imgs_list[idx].split('/')[-1].split('.jpg')[0]\n        rectangles = []\n\n        for box in self.df[self.df.image_id == id]['bbox'].values:\n            bbox = ast.literal_eval(box)\n            x = bbox[0]\n            y = bbox[1]\n            w = bbox[2]\n            h = bbox[3]\n            rectangles.append(patches.Rectangle((x,y),w,h,linewidth=1,edgecolor='r',facecolor='none'))\n\n        return rectangles\n\n    def format_boxes(self, boxes):\n        # replace width, height with xmax, ymax\n        try:\n            boxes[:, 2] =  boxes[:, 2] + boxes[:, 0]\n            boxes[:, 3] =  boxes[:, 3] + boxes[:, 1]\n        except:\n            pass\n        return boxes\n\n    def get_image(self, idx):\n        img_path = os.path.join(self.path, self.train_or_test_path, self.imgs_list[idx])\n        return np.array(Image.open(img_path).convert(\"RGB\"))\n\n    def draw(self, idx):\n        fig, ax = plt.subplots(1, figsize=(10, 10))\n        ax.imshow(dataset.get_image(idx))\n        for rectangle in dataset.get_rectangles(idx):\n            ax.add_patch(rectangle)\n        plt.show\n\n    def __getitem__(self, idx):\n        id = self.df.iloc[idx].image_id\n        boxes = np.int64(np.array([ast.literal_eval(box) for box in self.df[self.df.image_id == id]['bbox'].values]))\n\n        # format boxes width, height\n        boxes = self.format_boxes(boxes)\n        \n        target = {}\n        target[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.int64)\n        target[\"labels\"] = torch.ones((len(boxes),), dtype=torch.int64)\n        target[\"image_id\"] = torch.tensor([self.ids[id]])\n        target[\"area\"] = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        target[\"iscrowd\"] = torch.zeros((len(boxes),), dtype=torch.int64)\n        \n        img_path = os.path.join(self.path, self.train_or_test_path, self.imgs_list[idx])\n        img = Image.open(img_path).convert(\"RGB\")\n\n        if self.transforms is not None:\n            img = self.transforms(img)\n\n        return T.ToTensor()(img), target\n\n    def __len__(self):\n        return len(self.imgs_list)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transform(train):\n    transforms = []\n    if train:\n        # random horizontal flip with 50% probability\n        transforms.append(T.RandomHorizontalFlip(0.5))\n    return T.Compose(transforms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"path = '../input/global-wheat-detection/'\ntrain_path = 'train'\ntest_path = 'test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = GlobalWheatDataset(path, train_path, get_transform(train=True))\ndataset.draw(8)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Datasets and DataLoaders"},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def my_collate(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = GlobalWheatDataset(path, train_path, get_transform(train=True))\ndataset_test = GlobalWheatDataset(path, test_path, get_transform(train=False))\n\nindices = torch.randperm(len(dataset)).tolist()\nindices_test = torch.randperm(len(dataset_test)).tolist()\n\ndataset = torch.utils.data.Subset(dataset, indices)\ndataset_test = torch.utils.data.Subset(dataset_test, indices_test)\n\ndata_loader = torch.utils.data.DataLoader(\n        dataset, batch_size=2, shuffle=True, num_workers=4,\n        collate_fn=my_collate)\n\ndata_loader_test = torch.utils.data.DataLoader(\n        dataset_test, batch_size=1, shuffle=False, num_workers=4,\n        collate_fn=my_collate)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Faster RCNN Model with Resnet50"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n\nnum_classes = 2\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"params = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n\n# and a learning rate scheduler\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"total_errors = []\nfor epoch in range(10):\n    losses_arr = []\n\n    for images, targets in data_loader:\n\n        images = list(image.to(device) for image in images)\n        targets = [{k: torch.as_tensor(v).detach().to(device) for k, v in t.items()} for t in targets]\n\n        optimizer.zero_grad()\n\n        loss_dict = model(images, targets)\n        losses = sum(loss for loss in loss_dict.values())\n        losses_arr.append(losses.item())\n\n        losses.backward()\n        optimizer.step()\n\n        # update the learning rate\n        # lr_scheduler.step()\n        \n    total_errors.append(np.mean(np.array(losses_arr)))\n    if epoch % 1 == 0:\n        print(\"Epoch:{0:3d}, Loss:{1:1.3f}\".format(epoch, total_errors[-1]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Evaluate Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(total_errors)","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}