{"cells":[{"metadata":{},"cell_type":"markdown","source":"# AI - FasterRCNN Train\nhttps://www.kaggle.com/tomchaniii/pytorch-starter-fasterrcnn-inference\n"},{"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 re\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\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Download TorchVision repo to use some files from\n# references/detection\n!git clone https://github.com/pytorch/vision.git && cd vision && git checkout v0.3.0\n\n!cp vision/references/detection/utils.py ./\n!cp vision/references/detection/transforms.py ./\n!cp vision/references/detection/coco_eval.py ./\n!cp vision/references/detection/engine.py ./\n!cp vision/references/detection/coco_utils.py ./\n!rm -rf vision","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!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'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-500:]\ntrain_ids = image_ids[:-500]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_df.shape, train_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom PIL import Image\nimport math\n\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None, noise=0, size_ratio=1):\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        self.noise = noise\n        self.size_ratio = size_ratio\n\n    def __getitem__(self, index: int):\n        \n        if self.noise > 0:\n            is_noise = index % 2\n            index = math.floor(index / 2)\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n        image = Image.open(os.path.join(self.image_dir, f'{image_id}.jpg')).convert(\"RGB\")\n        image = image.resize((round(image.width * self.size_ratio), round(image.height * self.size_ratio)))\n        \n        if self.noise > 0:\n            if is_noise > 0:\n                noise_image = Image.effect_noise(image.size, 127).convert(\"RGB\")\n                image = Image.blend(image, noise_image, self.noise)\n                \n\n        boxes = records[['x', 'y', 'w', 'h']].values\n        \n        boxes[:, 2] = np.clip(np.round((boxes[:, 0] + boxes[:, 2]) * self.size_ratio), 0, image.width)\n        boxes[:, 3] = np.clip(np.round((boxes[:, 1] + boxes[:, 3]) * self.size_ratio), 0, image.height)\n        boxes[:, 0] = np.clip(np.round(boxes[:, 0] * self.size_ratio), 0, image.width)\n        boxes[:, 1] = np.clip(np.round(boxes[:, 1] * self.size_ratio), 0, image.height)\n        \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'] = torch.as_tensor(boxes, dtype=torch.int64)\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms is not None:\n            image, target = self.transforms(image, target)\n\n        return image, target\n\n    def __len__(self) -> int:\n        if self.noise > 0:\n            return self.image_ids.shape[0] * 2\n        else:\n            return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# helper functions for data augmentation / transformation\nfrom engine import train_one_epoch, evaluate\nimport transforms as T\n\n\ndef 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)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create the model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# load a model; pre-trained on COCO\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","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":"import utils\n\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, get_transform(train=True), noise=0.2, size_ratio=1)\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_transform(train=False))\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=16,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=utils.collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=16,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=utils.collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample, target = train_dataset[0]\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nsample = sample.permute(1,2,0).cpu().numpy()\nboxes = target['boxes'].data.cpu().numpy().astype(np.int32)\n\nsample = cv2.UMat(sample).get()\n\nfor box in boxes:\n    cv2.rectangle(\n        sample, \n        (box[0], box[1]), \n        (box[2], box[3]),\n        (220, 0, 0), 2)\n\nax.set_axis_off()\nax.imshow(sample)\n\nsample, target = train_dataset[1]\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nsample = sample.permute(1,2,0).cpu().numpy()\nboxes = target['boxes'].data.cpu().numpy().astype(np.int32)\n\nsample = cv2.UMat(sample).get()\n\nfor box in boxes:\n    cv2.rectangle(\n        sample, \n        (box[0], box[1]), \n        (box[2], box[3]),\n        (220, 0, 0), 2)\n\nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#for gpu\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\ntorch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.00349, momentum=0.9, weight_decay=0.0004)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = None\n\nnum_epochs = 8","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"/kaggle/working\")\n\nfrom engine import train_one_epoch, evaluate\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=20)\n    \n    # update the learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step()\n        \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":"# save model\ntorch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#remove downloaded modules\n!rm utils.py\n!rm transforms.py\n!rm coco_eval.py\n!rm engine.py\n!rm coco_utils.py","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}