{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Augmented prediction visualizer\n"},{"metadata":{},"cell_type":"markdown","source":"<img src=\"https://s5.gifyu.com/images/demo_trained_weights_1.gif\" align=\"center\"/>"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"%matplotlib notebook\nfrom typing import Optional, Dict, List\nfrom typing_extensions import TypedDict\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport pyarrow.parquet as pq\n\nimport matplotlib as mpl\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\n\nplt.style.use('seaborn')\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\n\nfrom albumentations import Compose, Normalize, PadIfNeeded\nfrom albumentations.pytorch import ToTensorV2\nfrom scipy.ndimage import zoom\nfrom IPython.display import HTML","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HEIGHT = 137\nWIDTH = 236\n\n# Preprocess\nTARGET_SIZE = 128\nPADDING = 8\n\nMEAN = 0.0778441\nSTD = 0.216016\n\n# You should keep this\nINPUT_PATH = '/kaggle/input/bengaliai-cv19'\n\n# Replace this to your weight dataset\nDATASET_PATH = '/kaggle/input/bengali-ai-model-weights'\n\n# Demo model's weights\nMODEL_STATE_FILE = DATASET_PATH + '/exp-549--resnet34--iter-24999.pt'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Preprocessing\nThe preprocessing script below is from Iafoss' kernel:\n\n[https://www.kaggle.com/iafoss/image-preprocessing-128x128](https://www.kaggle.com/iafoss/image-preprocessing-128x128)\n\nYou should apply the same preprocessing steps on the sample image as you did in your training process. See the example below."},{"metadata":{"trusted":true},"cell_type":"code","source":"def bbox(img):\n    rows = np.any(img, axis=1)\n    cols = np.any(img, axis=0)\n    rmin, rmax = np.where(rows)[0][[0, -1]]\n    cmin, cmax = np.where(cols)[0][[0, -1]]\n    return rmin, rmax, cmin, cmax\n\n\ndef crop_resize(img0, size=TARGET_SIZE, pad=64):\n    # crop a box around pixels large than the threshold\n    # some images contain line at the sides\n    ymin, ymax, xmin, xmax = bbox(img0[5:-5, 5:-5] > 80)\n\n    # cropping may cut too much, so we need to add it back\n    xmin = xmin - 13 if (xmin > 13) else 0\n    ymin = ymin - 10 if (ymin > 10) else 0\n    xmax = xmax + 13 if (xmax < WIDTH - 13) else WIDTH\n    ymax = ymax + 10 if (ymax < HEIGHT - 10) else HEIGHT\n    img = img0[ymin:ymax, xmin:xmax]\n\n    # remove lo intensity pixels as noise\n    img[img < 28] = 0\n    lx, ly = xmax - xmin, ymax - ymin\n    ls = max(lx, ly) + pad\n\n    # make sure that the aspect ratio is kept in rescaling\n    img = np.pad(img, [((ls - ly) // 2,), ((ls - lx) // 2,)], mode='constant')\n\n    return cv2.resize(img, (size, size))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Transformation functions\n\nSome transformation helper functions. You can implement your own:\n\n```python\ndef custom_transform(image, value):\n    \"\"\"Custom image transformation function.\n    \n    Args:\n        image (numpy.ndarray) Original image (without any previous transformation steps), shape: H x W\n        value (int|Tuple) Argument for the next transformation step\n        \n    Return:\n        transformed image (numpy.ndarray)\n    \"\"\"\n    transformed_image = ...\n\n    return transformed_image\n```"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rotate_image(image, angle):\n    \"\"\"Rotate.\"\"\"\n    image_center = tuple(np.array(image.shape[1::-1]) / 2)\n    matrix = cv2.getRotationMatrix2D(image_center, angle, 1.0)\n    result = cv2.warpAffine(image, matrix, image.shape[1::-1], flags=cv2.INTER_LINEAR,\n                            borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))\n\n    return result\n\ndef scale_image(image, scale):\n    \"\"\"Zoom in/out.\"\"\"\n    image_center = tuple(np.array(image.shape[1::-1]) / 2)\n    matrix = cv2.getRotationMatrix2D(image_center, 0, scale)\n    result = cv2.warpAffine(image, matrix, image.shape[1::-1], flags=cv2.INTER_LINEAR,\n                            borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))\n\n    return result\n\ndef get_range(start, minval, maxval, step=1.0):\n    \"\"\"Dummy step generator.\"\"\"\n    return [x for x in np.arange(start, maxval, step)] +\\\n           [x for x in np.arange(maxval, minval, -1 * step)] +\\\n           [x for x in np.arange(minval, start, step)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction Visualizer\nWith this simple class, you can visualize your model's output using different augmentations. It is a general script; you can use it (theoretically; I did not test it) for analyzing any image classification problem."},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelConfig(TypedDict):\n    label_id: str\n    label_name: str\n    num_classes: int\n\nclass Sample(TypedDict):\n    image: np.ndarray\n    image_id: str\n    image_mixup: Optional[np.ndarray]\n    labels: Dict[str, int]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"*The code is hidden for saving some space*"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class PredictionVisualizer:\n    \n    DARK_BLUE  = (0.1294, 0.5882, 0.9529) # True label (missclassified)\n    GREEN      = (0.2980, 0.6863, 0.3137) # Predicted: correct\n    RED        = (0.9869, 0.2627, 0.2118) # Predicted: incorrect\n    LIGHT_BLUE = (0.7333, 0.8706, 0.9843) # other\n    \n    def __init__(self):\n        \"\"\"Prediction visualizer.\n        \n        Usage:\n            See the example below in this notebook.\n        \"\"\"\n        self.__model = None\n        self.__labels = {}  # type: Dict[str, LabelConfig]\n        self.__sample = None  # type: Sample\n        self.__transform_fn = None  # type: Callable\n        self.__transform_fn_steps = None  # type: List[int]\n        self.__transofrm_data = []  # type: List[dict]\n    \n    @property\n    def model(self):\n        return self.__model\n    \n    @model.setter\n    def model(self, model):\n        self.__model = model\n\n    @property\n    def transform_fn(self):\n        return self.__transform_fn\n\n    def set_transform_fn(self, transform_fn, steps):\n        self.__transform_fn = transform_fn\n        self.__transform_fn_steps = steps\n    \n    @property\n    def sample(self) -> Sample:\n        return self.__sample\n\n    @sample.setter\n    def sample(self, sample: Sample):\n        self.__sample = sample\n    \n    def add_label(self, label_config: LabelConfig) -> None:\n        self.__labels[label_config['label_id']] = label_config\n    \n    @property\n    def labels(self) -> Dict[str, LabelConfig]:\n        return self.__labels\n    \n    def create(self):\n        \"\"\"Creates the animation.\"\"\"\n        fig, gs, subplots = self.__generate_figure()\n        frames = self.__generate_frames()\n        outputs = []\n        \n        for step_index, image in enumerate(frames['images']):\n            output = [\n                subplots['image'].imshow(image.astype(np.uint8), animated=True, cmap='Greys_r')\n            ]\n            \n            for label_idx, (label_id, label) in enumerate(self.__labels.items()):\n                colors = self.__get_colors(label['num_classes'],\n                                           self.sample['labels'][label_id],\n                                           np.argmax(frames[label_id], axis=1)[step_index]\n                                          )\n                output.append(\n                    subplots[label_id].vlines(np.array([x for x in range(0, label['num_classes'])]),\n                                              np.zeros(len(frames)),\n                                              frames[label_id][step_index],\n                                              colors\n                                             )\n                )\n                \n\n            outputs.append(output)\n\n        return animation.ArtistAnimation(fig, outputs, interval=50, blit=True, repeat=True, repeat_delay=2000)    \n    \n    def _before_forward(self, transformed_image):\n        \"\"\"Before forward adapter\n        \n        The default implementation converts the numpy array to torch tensor and adds the missing\n        `channel` and `batch` dimensions. You should update this method if your model expects a\n        different input format.\n\n        Args:\n            transformed_image (numpy.ndarray): Transformed image (rotated, etc), shape: H x W\n\n        Return:\n            Prepared images (batch of image 1) for your model. The returned image's shape should\n            be the shape of your model's input (For example, Pytorch: B x C x H x W)\n        \"\"\"        \n        # Convert to float tensor\n        transformed_image = torch.from_numpy(transformed_image).float()\n        \n        # Add 'channel' dim\n        transformed_image = transformed_image.unsqueeze(0)\n        \n        # Add 'batch' dimension\n        transformed_image = transformed_image.unsqueeze(0)\n        \n        return transformed_image\n    \n    def _forward(self, input_image):\n        \"\"\"You can make the forward call in here\n\n        Args: \n            input_image (torch.Tensor | any) Prepared input for your model. Shape: B x C x H x W\n        \n        Return:\n            You should return a dictionary of your model's predictions (logits or softmax)\n            for every registered labels.\n            \n            ```\n            with torch.no_grad():\n                out_graph, out_vowel, out_conso = self.model(input_image)\n            \n            return {\n                'grapheme_root': out_graph,\n                'vowel_diacritic': out_vowel,\n                'consonant_diacritic': out_conso\n            }\n            ```\n\n            out_x.shape => B x label.NUM_CLASS\n        \"\"\"\n        raise NotImplementedError\n\n    def _softmax(self, outputs):\n        \"\"\"Applies a softmax function and returns the result.\n\n        If your model has a final softmax layer, then you should override this to return\n        the `outputs` argument without changes.\n        \n        The visualizer will call this method for every label separately.\n        \n        Args:\n            outputs (torch.Tensor | any): Your model's output, shape: BATCH x NUM_CLASSES\n\n        Return:\n            Softmaxed values\n        \"\"\"\n        return F.softmax(outputs, dim=1)\n\n    def _after_forward(self, probabilities):\n        \"\"\"Convert the result to the required format.\n        \n        Args:\n            probabilities (torch.Tensor | any) Your model's output after the `self._softmax` call.\n            \n        Return: (numpy.ndarray)\n        \"\"\"\n        return probabilities.data.cpu().numpy()[0]\n    \n    def __generate_figure(self):\n        \"\"\"Generates the plot.\"\"\"\n        fig = plt.figure(constrained_layout=True, figsize=(14, 6))\n        gs = fig.add_gridspec(len(self.labels), 2)\n        \n        subplots = {}\n        subplots['image'] = fig.add_subplot(gs[:, 0], xticks=[], yticks=[])\n        subplots['image'].set_title('Image id: {}'.format(self.sample['image_id']), fontsize=10)\n\n        for label_idx, (label_id, label) in enumerate(self.__labels.items()):\n            subplots[label_id] = fig.add_subplot(gs[label_idx, 1], xlim=(-1, label['num_classes']))\n            subplots[label_id].set_title('{} (label: {})'.format(label['label_name'], self.sample['labels'][label_id]), fontsize=10)\n    \n        return fig, gs, subplots\n    \n    def __generate_frames(self):\n        \"\"\"Generates the frames.\"\"\"\n        \n        assert self.model is not None\n        assert self.sample is not None\n        assert self.transform_fn is not None\n        \n        h, w = self.sample['image'].shape\n        steps = len(self.__transform_fn_steps)\n        \n        frames = {}\n        \n        # Placeholder for the transformed images\n        frames['images'] = np.zeros((steps, h, w))\n        \n        # Create placeholders for the labels\n        for label_idx, (label_id, label) in enumerate(self.__labels.items()):\n            frames[label_id] = np.zeros((steps, label['num_classes']))\n            \n        for step, transform_step_value in enumerate(self.__transform_fn_steps):\n            \n            # Transform the original image\n            transformed_image = self.__transform_fn(self.sample['image'], transform_step_value)\n            \n            # Save the transformed image as a new frame\n            frames['images'][step, ...] = transformed_image\n            \n            # Prepare the image for the model\n            input_image = self._before_forward(transformed_image.copy())\n            \n            # Predict\n            model_output = self._forward(input_image)\n            \n            # Add the results to the frames\n            for label_id, output_logits in model_output.items():\n                frames[label_id][step, ...] = self._after_forward(self._softmax(output_logits))\n                \n        return frames\n\n    def __get_colors(self, size, target, pred):\n        \"\"\"Generates the colors of the vlines.\"\"\"\n        gra_color = [self.LIGHT_BLUE for _ in range(size)]\n\n        if pred == target:\n            gra_color[pred] = self.GREEN\n        else:\n            gra_color[pred] = self.RED\n            gra_color[target] = self.DARK_BLUE\n\n        return gra_color    \n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\nReplace this with your model. This is a Pytorch model, but you can use any other framework as well (I did not test it!)"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class BengaliModel(nn.Module):\n\n    def __init__(self, pretrained=False):\n        super().__init__()\n\n        self.backbone = torchvision.models.resnet34(pretrained=pretrained)\n        num_bottleneck_filters = self.backbone.fc.in_features\n        self.head_dropout = 0.1\n\n        self.backbone.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n\n        self.last = nn.Sequential(\n            nn.BatchNorm2d(num_bottleneck_filters),\n            nn.ReLU(inplace=True)\n        )\n\n        self.fc_graph = nn.Linear(num_bottleneck_filters, 168)\n        self.fc_vowel = nn.Linear(num_bottleneck_filters, 11)\n        self.fc_conso = nn.Linear(num_bottleneck_filters, 7)        \n\n    def forward(self, x):\n        batch_size, C, H, W = x.shape\n        \n        x = (x - MEAN * 255.0) / (STD * 255.0)\n\n        x = self.backbone.conv1(x)\n        x = self.backbone.bn1(x)\n        x = self.backbone.relu(x)\n        x = self.backbone.maxpool(x)\n\n        x = self.backbone.layer1(x)\n        x = self.backbone.layer2(x)\n        x = self.backbone.layer3(x)\n        x = self.backbone.layer4(x)\n\n        x = self.last(x)        \n\n        x = F.adaptive_avg_pool2d(x, 1).reshape(batch_size, -1)\n        x = F.dropout(x, self.head_dropout, self.training)\n\n        fc_graph = self.fc_graph(x)\n        fc_vowel = self.fc_vowel(x)\n        fc_conso = self.fc_conso(x)\n\n        return fc_graph, fc_vowel, fc_conso","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction analysis"},{"metadata":{},"cell_type":"markdown","source":"### Loading samples from parquet"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(INPUT_PATH + '/train.csv')\n\n# bengali_sample_id = 0\nbengali_sample_id = 13241\nbengali_sample = train_df.loc[bengali_sample_id].to_dict()\nbengali_sample","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"datafile = INPUT_PATH + '/train_image_data_{}.parquet'.format(0)\nparq = pq.read_pandas(datafile, columns=[str(x) for x in range(32332)]).to_pandas()\n\n# I trained my models using this inverted pixels, make sure this is matches with your training.\nparq = 255 - parq.iloc[:, :].values.reshape(-1, HEIGHT, WIDTH).astype(np.uint8)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Your custom visualizer\nTo make this work with your model, you have to update some of the visualizer's methods.\nThe default implementation works with Pytorch models, but you still may have to modify it."},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyVisualzier(PredictionVisualizer):\n    \n    def __init__(self):\n        super().__init__()\n\n    # The implementation of the `_forward` method is required.\n    # ---------------------------------------------------------\n    def _forward(self, input_image):\n\n        if torch.cuda.is_available():\n            input_image = input_image.cuda()\n        \n        with torch.no_grad():\n            out_graph, out_vowel, out_conso = self.model(input_image)\n\n        return {\n            'grapheme_root': out_graph,\n            'vowel_diacritic': out_vowel,\n            'consonant_diacritic': out_conso\n        }\n        \n        \n    # Implementation below this is optional\n    # ----------------------------------------------\n    def _before_forward(self, transformed_image):\n        \"\"\"Before forward adapter\n        \n        The default implementation converts the numpy array to torch tensor and adds the missing\n        `channel` and `batch` dimensions. You should update this method if your model expects a\n        different input format.\n\n        Args:\n            transformed_image (numpy.ndarray): Transformed image (rotated, etc), shape: H x W\n\n        Return:\n            Prepared images (a one element batch) for your model. The returned image's shape should\n            be the shape of your model's input (For example, Pytorch: B x C x H x W)\n        \"\"\"\n        return super()._before_forward(transformed_image)\n\n    def _softmax(self, outputs):\n        \"\"\"Applies a softmax function and returns the result.\n\n        If your model has a final softmax layer, then this method should return\n        the `outputs` argument without changes.\n        \n        The visualizer will call this method for every label separately.\n        \n        Args:\n            outputs (torch.Tensor | any): Your model's output, shape: BATCH x NUM_CLASSES\n\n        Return:\n            Softmaxed values\n        \"\"\"\n        return super()._softmax(outputs)\n    \n    def _after_forward(self, probabilities):\n        \"\"\"Convert the result to the required format.\n        \n        Args:\n            probabilities (torch.Tensor | any) Your model's output after the `self._softmax` call.\n            \n        Return: (numpy.ndarray)\n        \"\"\"\n        return super()._after_forward(probabilities)\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Testing with random weights"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = BengaliModel(pretrained=False)\n\n# If you'd like to use these demo weights, you have to\n# enable GPU first\n\n# My custom model state\nmodel_state = torch.load(MODEL_STATE_FILE)\n\n# Loading the weights\nmodel.load_state_dict(model_state['model_state_dict'])\n\n# Do not forget this!\nx = model.eval()\n\n# Move to the GPU if available\nif torch.cuda.is_available():\n    model.cuda()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"visualizer = MyVisualzier()\n\n# Set your model\nvisualizer.model = model\n\n# Add the labels to the visualizer\nvisualizer.add_label({'label_id': 'grapheme_root', 'label_name': 'Grapheme Root', 'num_classes': 168})\nvisualizer.add_label({'label_id': 'vowel_diacritic', 'label_name': 'Vowel Diacritic', 'num_classes': 11})\nvisualizer.add_label({'label_id': 'consonant_diacritic', 'label_name': 'Consonant Diacritic', 'num_classes': 7})\n\n# Create a new sample\nsample: Sample = {\n    # You should preprocess the image for your model\n    'image': crop_resize(parq[bengali_sample_id], size=TARGET_SIZE, pad=PADDING),\n    \n    # Image id\n    'image_id': 'Test_{}'.format(bengali_sample_id),\n\n    # True labels\n    'labels': {\n        'grapheme_root': bengali_sample['grapheme_root'],\n        'vowel_diacritic': bengali_sample['vowel_diacritic'],\n        'consonant_diacritic': bengali_sample['consonant_diacritic']\n    }\n}\n    \n# Set the new sample\nvisualizer.sample = sample\n\n# Set rotation.\nvisualizer.set_transform_fn(rotate_image, get_range(1, -45, 45, 1))\n# visualizer.set_transform_fn(scale_image, get_range(1, 0.25, 1.5, 0.02))\n\n# Create animation\nanim = visualizer.create()\n\n# Show the JS animation\nHTML(anim.to_jshtml())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# More examples\n(Animated gifs only)"},{"metadata":{},"cell_type":"markdown","source":"<img src=\"https://s5.gifyu.com/images/demo_trained_weights_1.gif\" width=\"600\" align=\"center\"/>"},{"metadata":{},"cell_type":"markdown","source":"<img src=\"https://s5.gifyu.com/images/demo_trained_weights_misscls_2.gif\" width=\"600\" align=\"center\"/>"},{"metadata":{},"cell_type":"markdown","source":"<img src=\"https://s5.gifyu.com/images/demo_trained_weights_misscls_scale_2.gif\" width=\"600\" align=\"center\"/>"},{"metadata":{},"cell_type":"markdown","source":"---------------\n**Thanks for reading.** If you find this notebook helpful (or interesting) please vote!"},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":1}