{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\nfrom torchvision.ops import box_iou, MultiScaleRoIAlign\nfrom torchvision import models\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\nfrom torch.utils.data import DataLoader, Dataset\nimport pytorch_lightning as pl\nfrom pytorch_lightning.metrics import AveragePrecision\n\nSEED=2484\nDEVICE=torch.device('cuda')\npl.utilities.seed.seed_everything(SEED)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class VBDDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transforms=None, phase='train'):\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.phase = phase\n\n    def __getitem__(self, idx):\n\n        image_id = self.image_ids[idx]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.png', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        if self.phase == 'test':\n            if self.transforms:\n                sample = {\n                    'image': image,\n                }\n                sample = self.transforms(**sample)\n                image = sample['image']\n            return image, image_id\n\n        boxes = records[['x_min', 'y_min', 'x_max', 'y_max']].values\n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n        # all the labels are shifted by 1 to accomodate background\n        labels = torch.squeeze(torch.as_tensor((records.class_id.values+1,), dtype=torch.int64))\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['area'] = area\n        target['image_id'] = torch.tensor([idx])\n        target['iscrowd'] = iscrowd\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.as_tensor(sample['bboxes'])\n\n        return image, target\n\n    def __len__(self):\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n#         A.Normalize(mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0, p=1.0),\n        ToTensorV2(p=1.0)\n    ])\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef format_prediction_string(labels, boxes, scores):\n    pred_strings = []\n    for j in zip(labels, scores, boxes):\n        pred_strings.append(\"{0} {1:.4f} {2} {3} {4} {5}\".format(\n            j[0], j[1], j[2][0], j[2][1], j[2][2], j[2][3]))\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class VinDetector(pl.LightningModule):\n    def __init__(self, **kwargs):\n        super().__init__()\n\n        self.model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n        num_classes = 15\n        in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n        self.learning_rate = 1e-3\n        self.batch_size = 4\n\n    def forward(self, x):\n        return self.model(x)\n\n    def prepare_data(self):\n        df = pd.read_csv('../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')\n        df = df[df['class_id'] != 14].reset_index(drop=True)\n        self.train_dataset = VBDDataset(df, '../input/vinbigdata-chest-xray-original-png/train', get_train_transform())\n\n    def train_dataloader(self):\n        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, pin_memory=True, num_workers=4, collate_fn=collate_fn)\n\n    def training_step(self, batch, batch_idx):\n        images, targets = batch\n        targets = [{k: v for k, v in t.items()} for t in targets]\n        loss_dict = self.model(images, targets)\n        loss = sum(loss for loss in loss_dict.values())\n        self.log('Loss', loss, on_step=True, on_epoch=True, prog_bar=True)\n        return {\"loss\": loss}\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate, momentum=0.95, weight_decay=1e-5, nesterov=True)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=6, eta_min=0, verbose=True)\n        return [optimizer], [scheduler]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net = VinDetector()\ntrainer = pl.Trainer(max_epochs=5, gpus=1, progress_bar_refresh_rate=100)\ntrainer.fit(net)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for p in net.model.parameters():\n    p.requires_grad = True\n\ntrainer = pl.Trainer(max_epochs=10, gpus=1, progress_bar_refresh_rate=100, precision=16)\ntrainer.fit(net)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('../input/vinbigdata-chest-xray-abnormalities-detection/sample_submission.csv')\ntest_dataset = VBDDataset(df, '../input/vinbigdata-chest-xray-original-png/test', get_valid_transform(), phase='test')\ntest_data_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=4, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"detection_threshold = 0.5\nresults = []\nnet.model.to(DEVICE)\nnet.model.eval()\n\nwith torch.no_grad():\n    for images, image_ids in test_data_loader:\n        images = list(image.to(DEVICE) for image in images)\n        outputs = net.model(images)\n        for i, image in enumerate(images):\n            image_id = image_ids[i]\n            result = {\n                'image_id': image_id,\n                'PredictionString': '14 1.0 0 0 1 1'\n            }\n            boxes = outputs[i]['boxes'].data.cpu().numpy()\n            labels = outputs[i]['labels'].data.cpu().numpy()\n            scores = outputs[i]['scores'].data.cpu().numpy()\n            if len(boxes) > 0:\n                labels = labels - 1\n                labels[labels == -1] = 14\n                selected = scores >= detection_threshold\n                boxes = boxes[selected].astype(np.int32)\n                scores = scores[selected]\n                labels = labels[selected]\n                if len(boxes) > 0:\n                    result = {\n                        'image_id': image_id,\n                        'PredictionString': format_prediction_string(labels, boxes, scores)\n                    }\n            results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('./fasterrcnn.csv', index=False)","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}