{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%%writefile model.py\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\n# https://www.kaggle.com/jonykarki/fasterrcnn-resnet101-training\ndef get_model():\n    backbone = resnet_fpn_backbone('resnet152', pretrained=True)\n    model = FasterRCNN(backbone, num_classes=2)\n    return model\n\nmodel = get_model()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"%%writefile train.py\n\nimport pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nimport albumentations as A\nimport torch\nimport torchvision\nimport time\n\nfrom model import model\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\nfrom PIL import Image\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom matplotlib import pyplot as plt\n# Albumentations\nfrom albumentations import (\n    HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose\n)\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nbatch_size = 8\n\nDIR_INPUT = '../input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\ntrain_df = pd.read_csv(f\"{DIR_INPUT}/train.csv\")\nprint(train_df.shape)\nprint(train_df.head())\n\ntrain_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)\n\nprint(train_df.head())\n\nimage_ids = train_df['image_id'].unique()\n# use almost all images for training\nvalid_ids = image_ids[-1:]\ntrain_ids = image_ids[:-1]\nprint(len(train_ids))\nprint(len(valid_ids))\n\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\n\nprint(valid_df.shape, train_df.shape)\n\nclass WheatDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_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: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n        \n#         image = Image.open(f\"{self.image_dir}/{image_id}.jpg\")\n#         image = image.convert('RGB')\n#         image = np.array(image)\n        image = cv2.imread(f\"{self.image_dir}/{image_id}.jpg\", cv2.IMREAD_COLOR).astype(np.float32)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n    \n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        # target['masks'] = None\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.FloatTensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\ndef get_train_transform():\n    # do a number of image augmentations\n    return A.Compose([\n        A.Flip(0.5),\n        A.RandomRotate90(0.5),\n        MotionBlur(p=0.2),\n        MedianBlur(blur_limit=3, p=0.1),\n        Blur(blur_limit=3, p=0.1),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\n\n# model\nmodel = model.to(device)\n\nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# dataloaders\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    # num_workers=4,\n    collate_fn=collate_fn\n)\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    # num_workers=4,\n    collate_fn=collate_fn\n)\n\nimages, targets, image_ids = next(iter(train_data_loader))\nimages = list(image.to(device) for image in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)\nlr_scheduler_increase = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=10.0)\nlr_scheduler_decrease = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n\nnum_epochs = 30\n\nloss_hist = Averager()\n\ndef fit():\n    itr = 1\n    model.train()\n    for epoch in range(num_epochs):\n        start = time.time()\n        \n        \n        loss_hist.reset()\n        \n        for images, targets, image_ids in train_data_loader:\n            \n            images = list(image.to(device) for image in images)\n            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n            \n            loss_dict = model(images, targets)\n\n            losses = sum(loss for loss in loss_dict.values())\n            loss_value = losses.item()\n\n            loss_hist.send(loss_value)\n\n            optimizer.zero_grad()\n            \n            losses.backward()\n            optimizer.step()\n\n            if itr % 50 == 0:\n                print(f\"Iteration #{itr} loss: {loss_value}\")\n\n            itr += 1\n\n        print(f\"Epoch #{epoch} loss: {loss_hist.value}\")   \n        end = time.time()\n        print(f\"Took {(end - start) / 60} minutes for epoch {epoch}\")\n\nfit()\n\n# save the model\ntorch.save(model.state_dict(), 'fasterrcnn_resnet152_fpn.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!python train.py","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}