{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\n\ntorch.__version__","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%%writefile models.py\nfrom torchvision.ops import MultiScaleRoIAlign\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision import models\n\n\n# Get the backbone of any pretrained network, we'll use AlexNet\nbackbone = models.mobilenet_v2(pretrained=True)\nnew_backbone = backbone.features\nnew_backbone.out_channels = 1280\n\n# Configure the anchors. We shall have 12 different anchors.\nnew_anchor_generator = AnchorGenerator(sizes=((128, 256, 512),), \n                                      aspect_ratios=((0.5, 1.0, 2.0),))\n\n# Configure the output size of ROI-Pooling layer. \n# We shall end up with (num_boxes, num_features, 4, 4) after the ROIPooling layer\nnew_roi_pooler = MultiScaleRoIAlign(featmap_names=['0'], output_size=4, sampling_ratio=1)\n\n\n# let's use dummy variables for mean, std, min_size and max_size\nmin_size = 300\nmax_size = 500\nmean = [0.485, 0.456, 0.406]\nstd = [0.229, 0.224, 0.225]\n\ndef get_model():\n    # Instantiate the Faster-rcnn model with the variables declared above.\n    model = FasterRCNN(backbone=new_backbone,\n                          num_classes=2, \n                          min_size=min_size, \n                          max_size=max_size, \n                          image_mean=mean, \n                          image_std=std, \n                          rpn_anchor_generator=new_anchor_generator, \n                          box_roi_pool=new_roi_pooler)\n    return model\nmodel = get_model()\n\"\"\"\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# def mobilenet_backbone(\n#     backbone_name,\n#     pretrained,\n#     fpn,\n#     norm_layer=misc_nn_ops.FrozenBatchNorm2d,\n#     trainable_layers=2,\n#     returned_layers=None,\n#     extra_blocks=None\n# ):\n\n# https://www.kaggle.com/jonykarki/fasterrcnn-resnet101-training\ndef get_model():\n    backbone = resnet_fpn_backbone('resnet18', pretrained=True)\n    model = FasterRCNN(backbone, num_classes=2)\n    return model\n\nmodel = get_model()\n\"\"\"","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 models 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_mobilenet_v2.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}