"""
[In the previous notebook](https://www.kaggle.com/purplejester/reading-tiff-images-with-python), I was trying to read the data into memory. But nothing was done about segmentation masks. Here I implement a simple way to decode the masks and show them on top of a randomly picked training image.
"""
from typing import Tuple

import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from osgeo import gdal
"""
## Reading TIFF Image
"""
def read_tiff(path: str, channels_last: bool = False) -> np.ndarray:
    """Reads TIFF file."""
    
    dataset = gdal.Open(path, gdal.GA_ReadOnly)
    n_channels = dataset.RasterCount
    width = dataset.RasterXSize
    height = dataset.RasterYSize
    image = np.zeros((n_channels, height, width), dtype=np.uint8)
    for i in range(n_channels):
        band = dataset.GetRasterBand(i+1)
        channel = band.ReadAsArray()
        image[i] = channel
    if channels_last:
        image = image.transpose(1, 2, 0)
    return image
"""
## Reading Segmentation Masks
"""
df_masks = pd.read_csv('/kaggle/input/hubmap-kidney-segmentation/train.csv')
df_masks.set_index('id', inplace=True)
df_masks
"""
As the dataset's descritpion reads, the masks are represented with running-length encoding. We convert these encodings into 2D masks to visualize segmentation.
"""
def rle_decode(encoded: str, shape: Tuple[int, int]):
    """Decodes an RLE-encoded string."""
    
    numbers = list(map(int, encoded.split()))
    starts, runs = [np.asarray(xs) for xs in (numbers[::2], numbers[1::2])]

    # pixels enumerations starts from 1 but arrays are 
    # indexed staring from 0 so need to make an adjustment
    starts -= 1  
    
    mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)
    for start, run in zip(starts, runs):
        mask[start:start + run] = 1
    
    # In NumPy arrays, first goes height and then goes width; also,
    # the pixels in the mask are enumerated from top to bottom and 
    # from left to right, but the mask was filled in a different 
    # order so need to transpose
    return mask.reshape(shape[1], shape[0]).T
"""
Before applying our decoder to the data, let's check if the implementation works correctly. We use a toy RLE-mask to see what the decoding function gives us back.
"""
n, m = 7, 9
encoded_mask = '9 1 16 4 24 1 39 1 46 1 51 5'
decoded_mask = rle_decode(encoded_mask, (n, m))

f, ax = plt.subplots(1, 1, figsize=(8, 6))
ax.imshow(decoded_mask, cmap='gray')
ax.axis('off')

for x in range(m):           # left to right
    for y in range(n):       # top to bottom
        index = x*n + y + 1  # pixels are enumerated starting from 1
        
        color = 'black' if decoded_mask[y][x] == 1 else 'white'
        
        ax.annotate(str(index), xy=(x, y), 
                    ha='center', va='center',
                    color=color, fontsize=18)
"""
We see that:
- `9 1` is mapped to a single 9th pixel,
- `16 4` is mapped to a vertical line in the 3rd column,
- other pixels filled accordingly to the mask as expected.

Therefore, we can conclude that our decoding method works as expected. Now let's try to apply it to the real data.


## Showing Segmentation Mask on Top of TIFF Images

As before, we read an image from the training subset and decode its segmentation mask.
"""
image_id = '0486052bb'
tiff_image = read_tiff(f'/kaggle/input/hubmap-kidney-segmentation/train/{image_id}.tiff')
tiff_mask = rle_decode(df_masks.loc[image_id].encoding, tiff_image.shape[1:])
tiff_image.shape, tiff_mask.shape
"""
The following function takes both the image and its segmentation mask, and combines them into a single image array. Note that the function makes some OpenCV transformations and resizing that can easilty fill up all the avaliable RAM in the kernel. Therefore, it can fail for too large images and/or too big `resize` parameter values.
"""
def overlay(
    image: np.ndarray,
    mask: np.ndarray,
    color: Tuple[int, int, int] = (255, 0, 0),
    alpha: float = 0.5, 
    resize: Tuple[int, int] = (1024, 1024)
) -> np.ndarray:
    """Combines image and its segmentation mask into a single image.
    
    Params:
        image: Training image.
        mask: Segmentation mask.
        color: Color for segmentation mask rendering.
        alpha: Segmentation mask's transparency.
        resize: If provided, both image and its mask are resized before blending them together.
    
    Returns:
        image_combined: The combined image.
        
    """
    color = np.asarray(color).reshape(3, 1, 1)
    colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
    masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
    image_overlay = masked.filled()
    
    if resize is not None:
        image = cv.resize(image.transpose(1, 2, 0), resize)
        image_overlay = cv.resize(image_overlay.transpose(1, 2, 0), resize)
    
    image_combined = cv.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)
    
    return image_combined
"""
Instead of trying to render the whole image, here we take some random patch and show its segments to speed the computations up. You can try to feed the full-size image. Especially, if running this code outside of Kaggle kernels.
"""
x0, y0 = 4096, 9000
x1, y1 = x0 + 1024, y0 + 1024
f, ax = plt.subplots(1, 1, figsize=(12, 12))
cut_overlay = overlay(tiff_image[:, y0:y1, x0:x1], tiff_mask[y0:y1, x0:x1], alpha=0.4)
ax.imshow(cut_overlay);
"""
Looking at the image, we can suppose that our code correctly reads both the data and masks and combines them together.
"""