{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torchvision\nfrom torch import nn, Tensor\nfrom torch.nn import functional as F\nfrom torch.jit.annotations import List, Tuple, Dict, Optional\n\n\nimport re\nimport pickle\nimport math\nimport numpy as np\n\nfrom torchvision import utils\nfrom PIL import Image\n\nimport os\nfrom os.path import join\nimport copy\n\nimport zlib\nimport base64\nimport albumentations as A\n\nimport random\nimport math\n\nfrom torchvision.models.detection.image_list import ImageList\nfrom torchvision.models.detection.roi_heads import paste_masks_in_image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install \"/kaggle/input/pycocotools/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl\"\n\nimport pycocotools\nfrom pycocotools import mask as cocomask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_classes = 20\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nfrom torchvision import models\n\ndef create_backbone(num_classes):\n    backbone_wide = resnet_fpn_backbone('resnet50', False)\n    conv1_weight = backbone_wide._modules['body']._modules['conv1'].weight.clone()\n    backbone_wide._modules['body']._modules['conv1'] = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)\n    with torch.no_grad():\n        backbone_wide._modules['body']._modules['conv1'].weight[:, :3] = conv1_weight\n        backbone_wide._modules['body']._modules['conv1'].weight[:, 3] = backbone_wide._modules['body']._modules['conv1'].weight[:, 0]\n    return backbone_wide","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model_instance_segmentation(num_classes):\n#     with open('/kaggle/input/resnet50pretraineddefault52percent-net/resnet50fpn_default_pretrained1.pkl', 'rb') as input:\n#         net = pickle.load(input)\n    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, \n                                                               progress=False, \n                                                               num_classes=num_classes, \n                                                               pretrained_backbone=False)\n\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n    hidden_layer = 256\n    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)\n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _resize_image_and_masks(image, self_min_size, self_max_size, target):\n    # type: (Tensor, float, float, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]\n    im_shape = torch.tensor(image.shape[-2:])\n    min_size = float(torch.min(im_shape))\n    max_size = float(torch.max(im_shape))\n    scale_factor = self_min_size / min_size\n    if max_size * scale_factor > self_max_size:\n        scale_factor = self_max_size / max_size\n        \n    image = torch.nn.functional.interpolate(\n        image[None], scale_factor=scale_factor, mode='bilinear', recompute_scale_factor=True,\n        align_corners=False)[0]\n\n    if target is None:\n        return image, target\n\n    if \"masks\" in target:\n        mask = target[\"masks\"]\n        \n        mask = F.interpolate(mask[:, None].float(), scale_factor=scale_factor)[:, 0].byte()\n        target[\"masks\"] = mask\n    return image, target\n\nclass GeneralizedRCNNTransform_(nn.Module):\n    \"\"\"\n    Performs input / target transformation before feeding the data to a GeneralizedRCNN\n    model.\n\n    The transformations it perform are:\n        - input normalization (mean subtraction and std division)\n        - input / target resizing to match min_size / max_size\n\n    It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets\n    \"\"\"\n\n    def __init__(self, min_size, max_size, image_mean, image_std):\n        super(GeneralizedRCNNTransform_, self).__init__()\n        if not isinstance(min_size, (list, tuple)):\n            min_size = (min_size,)\n        self.min_size = min_size\n        self.max_size = max_size\n        self.image_mean = image_mean\n        self.image_std = image_std\n\n    def forward(self,\n                images,       # type: List[Tensor]\n                targets=None  # type: Optional[List[Dict[str, Tensor]]]\n                ):\n        # type: (...) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]\n        images = [img for img in images]\n        if targets is not None:\n            # make a copy of targets to avoid modifying it in-place\n            # once torchscript supports dict comprehension\n            # this can be simplified as as follows\n            # targets = [{k: v for k,v in t.items()} for t in targets]\n            targets_copy: List[Dict[str, Tensor]] = []\n            for t in targets:\n                data: Dict[str, Tensor] = {}\n                for k, v in t.items():\n                    data[k] = v\n                targets_copy.append(data)\n            targets = targets_copy\n        for i in range(len(images)):\n            image = images[i]\n            target_index = targets[i] if targets is not None else None\n            \n#             print(\"target_index before = \", target_index)\n#             print(\"image.shape before = \", image.shape)\n\n#             if image.dim() != 3:\n#                 raise ValueError(\"images is expected to be a list of 3d tensors \"\n#                                  \"of shape [C, H, W], got {}\".format(image.shape))\n            image = self.normalize(image)\n            image, target_index = self.resize(image, target_index)\n            \n#             print(\"target_index after= \", target_index)\n#             print(\"image.shape after = \", image.shape)\n            \n            images[i] = image\n            if targets is not None and target_index is not None:\n                targets[i] = target_index\n\n        image_sizes = [img.shape[-2:] for img in images]\n        images = self.batch_images(images)\n        image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])\n        for image_size in image_sizes:\n            assert len(image_size) == 2\n            image_sizes_list.append((image_size[0], image_size[1]))\n\n        image_list = ImageList(images, image_sizes_list)\n    \n        return image_list, targets\n\n    def normalize(self, image):\n        dtype, device = image.dtype, image.device\n        mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)\n        std = torch.as_tensor(self.image_std, dtype=dtype, device=device)\n        return (image - mean[:, None, None]) / std[:, None, None]\n\n    def torch_choice(self, k):\n        # type: (List[int]) -> int\n        \"\"\"\n        Implements `random.choice` via torch ops so it can be compiled with\n        TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803\n        is fixed.\n        \"\"\"\n        index = int(torch.empty(1).uniform_(0., float(len(k))).item())\n        return k[index]\n\n    def resize(self, image, target):\n        # type: (Tensor, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]\n        h, w = image.shape[-2:]\n        if self.training:\n            size = float(self.torch_choice(self.min_size))\n        else:\n            # FIXME assume for now that testing uses the largest scale\n            size = float(self.min_size[-1])\n        if torchvision._is_tracing():\n            image, target = _resize_image_and_masks_onnx(image, size, float(self.max_size), target)\n        else:\n            image, target = _resize_image_and_masks(image, size, float(self.max_size), target)\n\n        if target is None:\n            return image, target\n\n        bbox = target[\"boxes\"]\n        bbox = resize_boxes(bbox, (h, w), image.shape[-2:])\n        target[\"boxes\"] = bbox\n\n        if \"keypoints\" in target:\n            keypoints = target[\"keypoints\"]\n            keypoints = resize_keypoints(keypoints, (h, w), image.shape[-2:])\n            target[\"keypoints\"] = keypoints\n        return image, target\n\n    # _onnx_batch_images() is an implementation of\n    # batch_images() that is supported by ONNX tracing.\n    @torch.jit.unused\n    def _onnx_batch_images(self, images, size_divisible=32):\n        # type: (List[Tensor], int) -> Tensor\n        max_size = []\n        for i in range(images[0].dim()):\n            max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)\n            max_size.append(max_size_i)\n        stride = size_divisible\n        max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)\n        max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)\n        max_size = tuple(max_size)\n\n        # work around for\n        # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)\n        # which is not yet supported in onnx\n        padded_imgs = []\n        for img in images:\n            padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]\n            padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))\n            padded_imgs.append(padded_img)\n\n        return torch.stack(padded_imgs)\n\n    def max_by_axis(self, the_list):\n        # type: (List[List[int]]) -> List[int]\n        maxes = the_list[0]\n        for sublist in the_list[1:]:\n            for index, item in enumerate(sublist):\n                maxes[index] = max(maxes[index], item)\n        return maxes\n\n    def batch_images(self, images, size_divisible=32):\n        # type: (List[Tensor], int) -> Tensor\n        if torchvision._is_tracing():\n            # batch_images() does not export well to ONNX\n            # call _onnx_batch_images() instead\n            return self._onnx_batch_images(images, size_divisible)\n\n        max_size = self.max_by_axis([list(img.shape) for img in images])\n        stride = float(size_divisible)\n        max_size = list(max_size)\n        max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)\n        max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)\n\n        batch_shape = [len(images)] + max_size\n        batched_imgs = images[0].new_full(batch_shape, 0)\n        for img, pad_img in zip(images, batched_imgs):\n            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)\n\n        return batched_imgs\n\n    def postprocess(self,\n                    result,               # type: List[Dict[str, Tensor]]\n                    image_shapes,         # type: List[Tuple[int, int]]\n                    original_image_sizes  # type: List[Tuple[int, int]]\n                    ):\n        # type: (...) -> List[Dict[str, Tensor]]\n        if self.training:\n            return result\n        for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):\n            boxes = pred[\"boxes\"]\n            boxes = resize_boxes(boxes, im_s, o_im_s)\n            result[i][\"boxes\"] = boxes\n            if \"masks\" in pred:\n                masks = pred[\"masks\"]\n                masks = paste_masks_in_image(masks, boxes, o_im_s)\n                result[i][\"masks\"] = masks\n            if \"keypoints\" in pred:\n                keypoints = pred[\"keypoints\"]\n                keypoints = resize_keypoints(keypoints, im_s, o_im_s)\n                result[i][\"keypoints\"] = keypoints\n        return result\n\n    def __repr__(self):\n        format_string = self.__class__.__name__ + '('\n        _indent = '\\n    '\n        format_string += \"{0}Normalize(mean={1}, std={2})\".format(_indent, self.image_mean, self.image_std)\n        format_string += \"{0}Resize(min_size={1}, max_size={2}, mode='bilinear')\".format(_indent, self.min_size,\n                                                                                         self.max_size)\n        format_string += '\\n)'\n        return format_string\n    \n    \ndef resize_boxes(boxes, original_size, new_size):\n    # type: (Tensor, List[int], List[int]) -> Tensor\n    ratios = [\n        torch.tensor(s, dtype=torch.float32, device=boxes.device) /\n        torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)\n        for s, s_orig in zip(new_size, original_size)\n    ]\n    ratio_height, ratio_width = ratios\n    xmin, ymin, xmax, ymax = boxes.unbind(1)\n\n    xmin = xmin * ratio_width\n    xmax = xmax * ratio_width\n    ymin = ymin * ratio_height\n    ymax = ymax * ratio_height\n    return torch.stack((xmin, ymin, xmax, ymax), dim=1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# means = {'red': 41.33620795268639, 'green': 25.602606457251067, 'blue': 27.628134909278703, 'yellow': 41.59902456019701}\n# stds = {'red': 68.24592496125408, 'green': 43.672423320044, 'blue': 74.81053636746947, 'yellow': 66.31413028835794}\n# counts = {'red': 515, 'green': 492, 'blue': 516, 'yellow': 501}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"net = get_model_instance_segmentation(num_classes)\nnet._modules['backbone'] = create_backbone(num_classes)\n\ntransform_ = GeneralizedRCNNTransform_((800,), 1333, [0, 0, 0, 0], [1, 1, 1, 1])\n# transform_ = GeneralizedRCNNTransform_((800,), 1333, [means['red']/counts['red'], means['green']/counts['green'], means['blue']/counts['blue'], means['yellow']/counts['yellow']], \n#                     [stds['red']/counts['red'], stds['green']/counts['green'], stds['blue']/counts['blue'], stds['yellow']/counts['yellow']])\n\nnet._modules['transform'] = transform_\n\nwith open('/kaggle/input/maskrcnn-full-dataset/maskrcnn_32percent_fulldataset4ch.pth', 'rb') as input:\n    model_weights = pickle.load(input)\n    net.load_state_dict(model_weights)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"net.to(device);","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from os import walk\n_, _, filenames = next(walk('/kaggle/input/hpa-single-cell-image-classification/test'))\n\nimg_ids = []\nfor file in filenames:\n    im_id = file.split('_')[0]\n    img_ids.append(im_id)\n\nunique_img_ids = set(img_ids)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from albumentations.pytorch.transforms import ToTensorV2\n\ntransforms_test = A.Compose([A.ToFloat(), ToTensorV2(transpose_mask=True, always_apply=True, p=1.0)])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HPADatasetTest(object):\n    def __init__(self, unique_img_ids, transforms):\n        self.transforms = transforms\n        self.img_list = list(unique_img_ids)\n\n    def __getitem__(self, idx):\n        img_dir = '/kaggle/input/hpa-single-cell-image-classification/test'\n        img_dict = {'red' : None, 'green' : None, 'blue' : None, 'yellow' : None}\n#         img_dict = {'red' : None, 'green' : None, 'blue' : None}\n        for key in img_dict.keys():\n            with Image.open(os.path.join(img_dir, self.img_list[idx] + \"_\" + key + \".png\")) as img:\n                img_dict[key] = np.array(img).astype(np.float32)\n#                 print(\"before\", idx, np.max(img_dict[key]))\n                if np.max(img_dict[key]) > 255:\n                      img_dict[key] = img_dict[key]*255/65535\n                img_dict[key] = np.floor(img_dict[key]).astype(np.ubyte)\n#                 print(\"after\", idx, np.max(img_dict[key]))\n                        \n#                 if np.max(img_dict[key]) > 1:\n#                       img_dict[key] = img_dict[key]/256\n        \n        img = np.stack((img_dict[k] for k in img_dict.keys()), axis = 2)\n        augmented = self.transforms(image = img)\n        img = augmented['image']\n        \n        return img, self.img_list[idx]\n\n    def __len__(self):\n        return len(self.img_list)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn_cust(batch):\n    return batch","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"testset = HPADatasetTest(unique_img_ids, transforms_test)\nbatch_size = 1\ntestloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, \n                                         num_workers=2, pin_memory=False, collate_fn=collate_fn_cust)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"min_score = 0.0\nf = open(\"submission.csv\", \"w\")\nf.write(\"ID,ImageWidth,ImageHeight,PredictionString\\n\")\nnet.eval()\nwith torch.no_grad():\n    for images in testloader:   \n        IDs = list(image[1] for image in images)\n        images = list(image[0].to(device) for image in images)\n        preds = net(images)\n        for pred_id, pred in enumerate(preds):\n            image_id = IDs[pred_id]\n            image_width = pred['masks'].shape[2]\n            image_height = pred['masks'].shape[3]\n            \n            line = image_id + ','+ str(image_width) + ',' + str(image_height) + ','\n            \n            for score_id, score in enumerate(pred['scores']):\n                if score > min_score:\n                    label = pred['labels'][score_id].item()\n                    if label == 19:\n                        label = 0\n                        \n                    mask = pred['masks'][score_id, 0, ::].cpu()\n                    mask = torch.where(mask>0.5, 1, 0)\n                    rle_mask = cocomask.encode(np.asfortranarray(mask.type(torch.uint8)))['counts']\n                    rle_mask_zlib_64 = base64.b64encode(zlib.compress(rle_mask))                    \n                    line += str(label) + ' ' + str(score.item()) + ' ' + rle_mask_zlib_64.decode('ascii') + ' '\n            line += '\\n'\n            f.write(line)\n            f.flush()\nf.close()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}