{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip3 install torch==0.4.1 albumentations\n\nimport glob\nimport math\nfrom multiprocessing.pool import ThreadPool\nimport time\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchvision\nfrom albumentations import (Compose, Normalize, Resize)\nfrom PIL import Image\nfrom scipy import ndimage\nfrom tqdm import tqdm\n\ndef get_images_in(path):\n    return np.sort(glob.glob(f'{path}/*.jpg'))\n\ndef read_image(path):\n    return np.array(Image.open(path))\n\ndef get_image_generator(image_paths, batch_size):\n    batch = []\n    for path in image_paths:\n        batch.append(read_image(path))\n        if len(batch) >= batch_size:\n            yield batch\n            batch = []\n    if len(batch) > 0: yield batch\n\ndef collate(batch):\n    if isinstance(batch[0], dict):\n        return {key: collate([sample[key] for sample in batch]) for key in batch[0].keys()}\n    return np.stack(batch)\n\ndef from_numpy(obj):\n    if isinstance(obj, dict):\n        return {key: from_numpy(value) for key, value in obj.items()}\n\n    if torch.cuda.is_available():\n        if isinstance(obj, torch.Tensor): return obj.float().cuda(non_blocking=True)\n        return torch.cuda.FloatTensor(obj)\n    else:\n        if isinstance(obj, torch.Tensor): return obj.float()\n        return torch.FloatTensor(obj)\n\ndef preprocess(pool, pipeline, batch):\n    return from_numpy(collate(list(pool.map(pipeline, batch))))\n\ndef channels_first(image):\n    return np.moveaxis(image, 2, 0)\n\nclass ChannelsFirst:\n    def __call__(self, **args):\n        args['image'] = channels_first(args['image'])\n        return args\n\ndef classifier_pipeline(image):\n    return Compose([\n        Resize(224, 224),\n        ChannelsFirst()\n    ])(image=image)\n\ndef segmenter_pipeline(image):\n    return Compose([\n        ChannelsFirst()\n    ])(image=image)\n\ndef as_cuda(tensor):\n    if torch.cuda.is_available():\n        return tensor.cuda()\n    return tensor\n\ndef extract_instance_masks_from_binary_mask(args):\n    _id, binary_mask = args\n    masks = []\n    labelled_mask = ndimage.label(binary_mask)[0]\n    labels, areas = np.unique(labelled_mask, return_counts=True)\n    labels = labels[areas >= 80]\n    for label in labels:\n        if label == 0: continue\n        masks.append((_id, labelled_mask == label))\n    if len(masks) < 1: return [(_id, None)]\n    return masks\n\ndef encode_rle(args):\n    _id, mask = args\n    if mask is None: return (_id, None)\n    pixels = mask.T.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return (_id, ' '.join(str(x) for x in runs))\n\ndef postprocess_segmentation(pool, ids, binary_masks):\n    ids_and_instance_masks = map(extract_instance_masks_from_binary_mask, zip(ids, binary_masks))\n    return map(encode_rle, sum(ids_and_instance_masks, []))\n\nclass Decoder(torch.nn.Module):\n    def __init__(self, in_channels, out_channels, stride, output_padding):\n        super().__init__()\n        self.layers = torch.nn.Sequential(\n            torch.nn.Conv2d(in_channels, in_channels // 4, (1, 1)),\n            torch.nn.BatchNorm2d(in_channels // 4),\n            torch.nn.ReLU(inplace=True),\n            torch.nn.ConvTranspose2d(in_channels // 4, in_channels // 4, (3, 3), stride=stride, padding=1, output_padding=output_padding),\n            torch.nn.BatchNorm2d(in_channels // 4),\n            torch.nn.ReLU(inplace=True),\n            torch.nn.Conv2d(in_channels // 4, out_channels, (1, 1)),\n            torch.nn.BatchNorm2d(out_channels),\n            torch.nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.layers(x)\n\nclass Segmenter(torch.nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.resnet = torchvision.models.resnet18(pretrained=True)\n        self.decoder1 = Decoder(512, 256, stride=2, output_padding=1)\n        self.decoder2 = Decoder(256, 128, stride=2, output_padding=1)\n        self.decoder3 = Decoder(128, 64, stride=2, output_padding=1)\n        self.decoder4 = Decoder(64, 64, stride=1, output_padding=0)\n        self.classifier = torch.nn.Sequential(\n            torch.nn.ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=2, padding=1, output_padding=1),\n            torch.nn.BatchNorm2d(32),\n            torch.nn.ReLU(inplace=True),\n            torch.nn.Conv2d(32, 32, kernel_size=(3, 3), stride=1, padding=1),\n            torch.nn.BatchNorm2d(32),\n            torch.nn.ReLU(inplace=True),\n            torch.nn.ConvTranspose2d(32, num_classes, (2, 2), stride=2)\n        )\n\n    def forward(self, x):\n        x = x['image']\n        x = self.resnet.conv1(x)\n        x = self.resnet.bn1(x)\n        x = self.resnet.relu(x)\n        x = self.resnet.maxpool(x)\n        x1 = self.resnet.layer1(x)\n        x2 = self.resnet.layer2(x1)\n        x3 = self.resnet.layer3(x2)\n        x4 = self.resnet.layer4(x3)\n        x = self.decoder1(x4) + x3\n        x = self.decoder2(x) + x2\n        x = self.decoder3(x) + x1\n        x = self.decoder4(x)\n        return {'mask': self.classifier(x)}\n\nclass Classifier(torch.nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.resnet = torchvision.models.resnet18(pretrained=True)\n        self.resnet.fc = torch.nn.Linear(self.resnet.fc.in_features, num_classes)\n\n    def forward(self, x):\n        x = x['image']\n        return {'has_ships': self.resnet(x)}\n\nclassifier_path = '/kaggle/input/airbus-model-weights/classifier_weights.pt'\nsegmenter_path = '/kaggle/input/airbus-model-weights/segmenter_weights.pt'\noutput_path = '/kaggle/working/submission.csv'\ndirectory_path = '/kaggle/input/airbus-ship-detection/test_v2'\nclassifier_batch_size = 16\nsegmenter_batch_size = 16\n\ntorch.backends.cudnn.benchmark = True\ntorch.set_grad_enabled(False)\n\nclassifier = Classifier(1).cuda()\nclassifier.load_state_dict(torch.load(classifier_path))\nclassifier.eval()\n\nsegmenter = Segmenter(1).cuda()\nsegmenter.load_state_dict(torch.load(segmenter_path))\nsegmenter.eval()\n\n# Warmup\nfor i in range(10):\n    batch = torch.randn(classifier_batch_size, 3, 224, 224).cuda()\n    classifier({'image': batch})\n\nfor i in range(10):\n    batch = torch.randn(segmenter_batch_size, 3, 768, 768).cuda()\n    segmenter({'image': batch})\n    \nimage_paths = get_images_in(directory_path)\n\ntotal_inference_time = 0\n\nmeans = from_numpy(np.array([0.485, 0.456, 0.406])[None, :, None, None])\nstds = from_numpy(np.array([0.229, 0.224, 0.225])[None, :, None, None])\n\npool = ThreadPool(2)\n\n# Classification\npred_labels = []\nfor batch in tqdm(get_image_generator(image_paths, classifier_batch_size), total=math.ceil(len(image_paths) / classifier_batch_size)):\n    time_start = time.time()\n    batch = preprocess(pool, classifier_pipeline, batch)\n    batch['image'] = (batch['image'] / 255 - means) / stds\n    pred_labels.extend((classifier(from_numpy(batch))['has_ships'] > 0)[:, 0])\n    torch.cuda.synchronize()\n    total_inference_time += (time.time() - time_start)\n\ntime_start = time.time()\npred_labels = np.array(pred_labels)\npositive_image_paths = image_paths[pred_labels == 1]\nnegative_image_paths = image_paths[pred_labels == 0]\ntotal_inference_time += (time.time() - time_start)\n\n# Segmentation\ntime_start = time.time()\nrecords = []\nremaining_ids = list(map(lambda path: path.split('/')[-1], positive_image_paths))\ntotal_inference_time += (time.time() - time_start)\nfor batch in tqdm(get_image_generator(positive_image_paths, segmenter_batch_size), total=math.ceil(len(positive_image_paths) / segmenter_batch_size)):\n    time_start = time.time()\n    batch = preprocess(pool, segmenter_pipeline, batch)\n    batch['image'] = (batch['image'] / 255 - means) / stds\n    binary_masks = (segmenter(batch)['mask'] > 0)[:, 0]\n    records.extend(postprocess_segmentation(pool, remaining_ids[:len(binary_masks)], binary_masks))\n    remaining_ids = remaining_ids[len(binary_masks):]\n    torch.cuda.synchronize()\n    total_inference_time += (time.time() - time_start)\n\ntime_start = time.time()\nnegative_ids = list(map(lambda path: path.split('/')[-1], negative_image_paths))\nrecords.extend(map(lambda _id: (_id, None), negative_ids))\n\nimage_ids, encoded_pixels = zip(*records)\ndf = pd.DataFrame({'ImageId': image_ids, 'EncodedPixels': encoded_pixels})\ndf.to_csv(output_path, index=False)\ntotal_inference_time += (time.time() - time_start)\nprint('Inference Time: %0.2f Minutes'%((total_inference_time)/60))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","collapsed":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}