<filename>src/utils/mask_utils.py
from PIL import Image
import cv2
import matplotlib.pyplot as plt
import pandas as pd
from src.utils import rle_utils as rle
import numpy as np

# Code modified from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def get_img_and_mask(img_path, annotation, width, height, mask_only=False):
    """ Capture the relevant image array as well as the image mask """
    img_mask = np.zeros((height, width), dtype=np.int16)
    for i, annot in enumerate(annotation): 
        img_mask = np.where(rle.decode(annot, (height, width))!=0, i, img_mask)
    
    if mask_only:
        return img_mask
    
    img = np.array(Image.open(img_path), dtype=np.int16)

    return img, img_mask

# Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def plot_img_and_mask(img, mask):
    """ Function to take an image and the corresponding mask and plot
    
    Args:
        img (np.arr): 1 channel np arr representing the image of cellular structures
        mask (np.arr): 1 channel np arr representing the instance masks (incrementing by one)
        
    Returns:
        None; Plots the two arrays and overlays them to create a merged image
    """
    plt.figure(figsize=(20,10))
    
    plt.subplot(1,3,1)
    _img = np.tile(np.expand_dims(img, axis=-1), 3)
    plt.imshow(_img)
    plt.axis(False)
    plt.title("Cell Image", fontweight="bold")
    
    plt.subplot(1,3,2)
    _mask = np.zeros_like(_img)
    _mask[..., 0] = mask
    plt.imshow(mask, cmap="inferno")
    plt.axis(False)
    plt.title("Instance Segmentation Mask", fontweight="bold")
    
    merged = cv2.addWeighted(_img, 0.75, np.clip(_mask, 0, 1)*255, 0.25, 0.0,)
    plt.subplot(1,3,3)
    plt.imshow(merged)
    plt.axis(False)
    plt.title("Cell Image w/ Instance Segmentation Mask Overlay", fontweight="bold")
    
    plt.tight_layout()
    plt.show()


# Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def save_mask(row, path = ''):
    msk = get_img_and_mask(**row[["img_path", "annotation", "width", "height"]].to_dict(), mask_only=True)
    np.savez(path+"{}".format(row.id), msk)

<filename>src/utils/mask_utils.py
from PIL import Image
import cv2
import matplotlib.pyplot as plt
import pandas as pd
from src.utils import rle_utils as rle
import numpy as np

# Code modified from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def get_img_and_mask(img_path, annotation, width, height, mask_only=False):
    """ Capture the relevant image array as well as the image mask """
    img_mask = np.zeros((height, width), dtype=np.int16)
    for i, annot in enumerate(annotation): 
        img_mask = np.where(rle.decode(annot, (height, width))!=0, i, img_mask)
    
    if mask_only:
        return img_mask
    
    img = np.array(Image.open(img_path), dtype=np.int16)

    return img, img_mask

# Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def plot_img_and_mask(img, mask):
    """ Function to take an image and the corresponding mask and plot
    
    Args:
        img (np.arr): 1 channel np arr representing the image of cellular structures
        mask (np.arr): 1 channel np arr representing the instance masks (incrementing by one)
        
    Returns:
        None; Plots the two arrays and overlays them to create a merged image
    """
    plt.figure(figsize=(20,10))
    
    plt.subplot(1,3,1)
    _img = np.tile(np.expand_dims(img, axis=-1), 3)
    plt.imshow(_img)
    plt.axis(False)
    plt.title("Cell Image", fontweight="bold")
    
    plt.subplot(1,3,2)
    _mask = np.zeros_like(_img)
    _mask[..., 0] = mask
    plt.imshow(mask, cmap="inferno")
    plt.axis(False)
    plt.title("Instance Segmentation Mask", fontweight="bold")
    
    merged = cv2.addWeighted(_img, 0.75, np.clip(_mask, 0, 1)*255, 0.25, 0.0,)
    plt.subplot(1,3,3)
    plt.imshow(merged)
    plt.axis(False)
    plt.title("Cell Image w/ Instance Segmentation Mask Overlay", fontweight="bold")
    
    plt.tight_layout()
    plt.show()


# Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset
def save_mask(row, path = ''):
    msk = get_img_and_mask(**row[["img_path", "annotation", "width", "height"]].to_dict(), mask_only=True)
    np.savez(path+"{}".format(row.id), msk)

# Code modified from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset Capture the relevant image array as well as the image mask # Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset Function to take an image and the corresponding mask and plot          Args:         img (np.arr): 1 channel np arr representing the image of cellular structures         mask (np.arr): 1 channel np arr representing the instance masks (incrementing by one)              Returns:         None; Plots the two arrays and overlays them to create a merged image # Code from https://www.kaggle.com/dschettler8845/sartorius-segmentation-mask-dataset#create_dataset