{"cells":[{"metadata":{},"cell_type":"markdown","source":"# WBF approach over TTA for single model Resnest Fasterrcnn\n\nHi everyone!\n\nThis is my first open kernel which use TTA for resnest fastrrcnn.\n\nMy resnest fasterrcnn is based on torchvision's fasterrcnn.\n\nIf it is helpful, please upvote my kernel","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Main Idea\n\nThis kernel is based on https://www.kaggle.com/shonenkov/wbf-over-tta-single-model-efficientdet","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Dependencies","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":true,"_kg_hide-input":true},"cell_type":"code","source":"!pip install --no-deps '../input/timm-package/timm-0.1.26-py3-none-any.whl' > /dev/null\n!pip install --no-deps '../input/pycocotools/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl' > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"../input/omegaconf\")\nsys.path.insert(0, \"../input/weightedboxesfusion\")\nsys.path.insert(0, \"../input/resnest/pytorch-image-models-master/pytorch-image-models-master\")\n\nimport ensemble_boxes\nimport torch\nimport torchvision\nfrom  torchvision.models.utils import load_state_dict_from_url\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torchvision.models.detection.backbone_utils import BackboneWithFPN\nfrom torchvision.ops import misc as misc_nn_ops\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\nfrom matplotlib import pyplot as plt\nfrom collections import OrderedDict\nfrom torch import nn\nimport warnings\nfrom torch.jit.annotations import Tuple, List, Dict, Optional\nfrom timm.models.resnest import resnest101e\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nfrom torch.utils.data import Dataset,DataLoader\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport cv2\nimport gc\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\nDIR_WEIGHTS = '/kaggle/input/resnest'\n\nWEIGHTS_FILE = f'{DIR_WEIGHTS}/best-model-sgd-cos.bin'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_transforms():\n    return A.Compose([\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT_PATH = '../input/global-wheat-detection/test'\n\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, image_ids, transforms=None):\n        super().__init__()\n        self.image_ids = image_ids\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{DATA_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        if self.transforms:\n            sample = {'image': image}\n            sample = self.transforms(**sample)\n            image = sample['image']\n        return image, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = DatasetRetriever(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{DATA_ROOT_PATH}/*.jpg')]),\n    transforms=get_valid_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndata_loader = DataLoader(\n    dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=2,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CrossEntropyLabelSmooth(nn.Module):\n    \"\"\"Cross entropy loss with label smoothing regularizer.\n\n    Reference:\n    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.\n    Equation: y = (1 - epsilon) * y + epsilon / K.\n\n    Args:\n        num_classes (int): number of classes.\n        epsilon (float): weight.\n    \"\"\"\n    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):\n        super(CrossEntropyLabelSmooth, self).__init__()\n        self.num_classes = num_classes\n        self.epsilon = epsilon\n        self.use_gpu = use_gpu\n        self.logsoftmax = nn.LogSoftmax(dim=1)\n\n    def forward(self, inputs, targets):\n        \"\"\"\n        Args:\n            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)\n            targets: ground truth labels with shape (num_classes)\n        \"\"\"\n        log_probs = self.logsoftmax(inputs)\n        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)\n        if self.use_gpu: targets = targets.cuda()\n        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes\n        loss = (- targets * log_probs).mean(0).sum()\n        return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fastrcnn_loss(class_logits, box_regression, labels, regression_targets):\n    # type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor]\n    \"\"\"\n    Computes the loss for Faster R-CNN.\n    Arguments:\n        class_logits (Tensor)\n        box_regression (Tensor)\n        labels (list[BoxList])\n        regression_targets (Tensor)\n    Returns:\n        classification_loss (Tensor)\n        box_loss (Tensor)\n    \"\"\"\n\n    labels = torch.cat(labels, dim=0)\n    regression_targets = torch.cat(regression_targets, dim=0)\n    labal_smooth_loss = CrossEntropyLabelSmooth(2)\n    classification_loss = labal_smooth_loss(class_logits, labels)\n\n    # get indices that correspond to the regression targets for\n    # the corresponding ground truth labels, to be used with\n    # advanced indexing\n    sampled_pos_inds_subset = torch.nonzero(labels > 0).squeeze(1)\n    labels_pos = labels[sampled_pos_inds_subset]\n    N, num_classes = class_logits.shape\n    box_regression = box_regression.reshape(N, -1, 4)\n\n    box_loss = det_utils.smooth_l1_loss(\n        box_regression[sampled_pos_inds_subset, labels_pos],\n        regression_targets[sampled_pos_inds_subset],\n        beta=1 / 9,\n        size_average=False,\n    )\n    box_loss = box_loss / labels.numel()\n\n    return classification_loss, box_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def resnest_fpn_backbone(pretrained, norm_layer=misc_nn_ops.FrozenBatchNorm2d, trainable_layers=3):\n    # resnet_backbone = resnet.__dict__['resnet152'](pretrained=pretrained,norm_layer=norm_layer)\n    backbone = resnest101e(pretrained=pretrained)\n    # select layers that wont be frozen\n    assert trainable_layers <= 5 and trainable_layers >= 0\n    layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]\n    # freeze layers only if pretrained backbone is used\n    for name, parameter in backbone.named_parameters():\n        if all([not name.startswith(layer) for layer in layers_to_train]):\n            parameter.requires_grad_(False)\n    return_layers = {'layer1': '0', 'layer2': '1', 'layer3': '2', 'layer4': '3'}\n    in_channels_stage2 = backbone.inplanes // 8\n    in_channels_list = [\n        in_channels_stage2,\n        in_channels_stage2 * 2,\n        in_channels_stage2 * 4,\n        in_channels_stage2 * 8,\n    ]\n    out_channels = 256\n    return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDetector(torch.nn.Module):\n    def __init__(self, trainable_layers=3, **kwargs):\n        super(WheatDetector, self).__init__()\n        backbone = resnest_fpn_backbone(pretrained=False)\n        self.base = FasterRCNN(backbone, num_classes = 2, **kwargs)\n        self.base.roi_heads.fastrcnn_loss = fastrcnn_loss\n\n    def forward(self, images, targets=None):\n        return self.base(images, targets)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_net(checkpoint_path):\n    model = WheatDetector()\n\n    # Load the trained weights\n    checkpoint=torch.load(checkpoint_path)\n    model.load_state_dict(checkpoint['model_state_dict'])\n\n\n    del checkpoint\n    gc.collect()\n\n    model.eval();\n    return model.cuda()\n\nnet = load_net(WEIGHTS_FILE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom TTA API\n\nIdea is simple: \n- `augment` make tta for one image\n- `batch_augment` make tta for batch of images\n- `deaugment_boxes` return tta predicted boxes in back to original state of image\n\nAlso we are interested in `Compose` with combinations of tta :)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 1024\n\n    def augment(self, image):\n        raise NotImplementedError\n    \n    def batch_augment(self, images):\n        raise NotImplementedError\n    \n    def deaugment_boxes(self, boxes):\n        raise NotImplementedError\n\nclass TTAHorizontalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n\n    def augment(self, image):\n        return image.flip(1)\n    \n    def batch_augment(self, images):\n        return images.flip(2)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [1,3]] = self.image_size - boxes[:, [3,1]]\n        return boxes\n\nclass TTAVerticalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return image.flip(2)\n    \n    def batch_augment(self, images):\n        return images.flip(3)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]\n        return boxes\n    \nclass TTARotate90(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 1, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 1, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = self.image_size - boxes[:, [3,1]] \n        res_boxes[:, [1,3]] = boxes[:, [0,2]]\n        return res_boxes\n\nclass TTACompose(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def augment(self, image):\n        for transform in self.transforms:\n            image = transform.augment(image)\n        return image\n    \n    def batch_augment(self, images):\n        for transform in self.transforms:\n            images = transform.batch_augment(images)\n        return images\n    \n    def prepare_boxes(self, boxes):\n        result_boxes = boxes.copy()\n        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)\n        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)\n        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)\n        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)\n        return result_boxes\n    \n    def deaugment_boxes(self, boxes):\n        for transform in self.transforms[::-1]:\n            boxes = transform.deaugment_boxes(boxes)\n        return self.prepare_boxes(boxes)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Demonstration how it works:","execution_count":null},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"def process_det(index, outputs, score_threshold=0.5):\n    boxes = outputs[index]['boxes'].data.cpu().numpy()   \n    scores = outputs[index]['scores'].data.cpu().numpy()\n    boxes = (boxes).clip(min=0, max=1023).astype(int)\n    indexes = np.where(scores>score_threshold)\n    boxes = boxes[indexes]\n    scores = scores[indexes]\n    return boxes, scores","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# you can try own combinations:\ntransform = TTACompose([\n    TTARotate90(),\n    TTAVerticalFlip(),\n    TTAHorizontalFlip(),\n])\n\nfig, ax = plt.subplots(1, 3, figsize=(16, 6))\n\nimage, image_id = dataset[5]\n\nnumpy_image = image.permute(1,2,0).cpu().numpy().copy()\n\nax[0].imshow(numpy_image);\nax[0].set_title('original')\n\ntta_image = transform.augment(image)\ntta_image_numpy = tta_image.permute(1,2,0).cpu().numpy().copy()\n\noutputs = net(tta_image.unsqueeze(0).float().cuda())\nboxes, scores = process_det(0, outputs)\n\nfor box in boxes:\n    cv2.rectangle(tta_image_numpy, (box[0], box[1]), (box[2],  box[3]), (0, 1, 0), 2)\n\nax[1].imshow(tta_image_numpy);\nax[1].set_title('tta')\n    \nboxes = transform.deaugment_boxes(boxes)\n\nfor box in boxes:\n    cv2.rectangle(numpy_image, (box[0], box[1]), (box[2],  box[3]), (0, 1, 0), 2)\n    \nax[2].imshow(numpy_image);\nax[2].set_title('deaugment predictions');","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Combinations of TTA","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from itertools import product\n\ntta_transforms = []\nfor tta_combination in product([TTAHorizontalFlip(), None], \n                               [TTAVerticalFlip(), None],\n                               [TTARotate90(), None]):\n    tta_transforms.append(TTACompose([tta_transform for tta_transform in tta_combination if tta_transform]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# WBF over TTA:","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_tta_predictions(images, score_threshold=0.5):\n    with torch.no_grad():\n        images = torch.stack(images).float().cuda()\n        predictions = []\n        for tta_transform in tta_transforms:\n            result = []\n            outputs = net(tta_transform.batch_augment(images.clone()))\n\n            for i, image in enumerate(images):\n                boxes = outputs[i]['boxes'].data.cpu().numpy()   \n                scores = outputs[i]['scores'].data.cpu().numpy()\n                indexes = np.where(scores > score_threshold)[0]\n                boxes = boxes[indexes]\n                boxes = tta_transform.deaugment_boxes(boxes.copy())\n                result.append({\n                    'boxes': boxes,\n                    'scores': scores[indexes],\n                })\n            predictions.append(result)\n    return predictions\n\ndef run_wbf(predictions, image_index, image_size=1024, iou_thr=0.5, skip_box_thr=0.43, weights=None):\n    boxes = [(prediction[image_index]['boxes']/(image_size-1)).tolist() for prediction in predictions]\n    scores = [prediction[image_index]['scores'].tolist() for prediction in predictions]\n    labels = [np.ones(prediction[image_index]['scores'].shape[0]).astype(int).tolist() for prediction in predictions]\n    boxes, scores, labels = ensemble_boxes.ensemble_boxes_wbf.weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n    boxes = boxes*(image_size-1)\n    return boxes, scores, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nfor j, (images, image_ids) in enumerate(data_loader):\n    break\n\npredictions = make_tta_predictions(images)\n\ni = 1\nsample = images[i].permute(1,2,0).cpu().numpy()\n\nboxes, scores, labels = run_wbf(predictions, image_index=i)\nboxes = boxes.round().astype(np.int32).clip(min=0, max=1023)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (1, 0, 0), 1)\n\nax.set_axis_off()\nax.imshow(sample);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = []\n\nfor images, image_ids in data_loader:\n    predictions = make_tta_predictions(images)\n    for i, image in enumerate(images):\n        boxes, scores, labels = run_wbf(predictions, image_index=i)\n        boxes = boxes.round().astype(np.int32).clip(min=0, max=1023)\n        image_id = image_ids[i]\n\n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n\n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(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.to_csv('submission.csv', index=False)\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Thank you for reading my kernel!","execution_count":null}],"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}