Tools.py

#!/usr/bin/env python
# coding: utf-8

# In[ ]:


# Data Science
import numpy  as np
import pandas as pd
# Visualization
import seaborn as sns
import matplotlib.pyplot as plt
# Tricks
sns.set(style='ticks', context='talk', font_scale=1.15)


# In[ ]:


import os, sys
from skimage.io import imread as skIR
from PIL import Image

# Root directory of the project
ROOT_DIR = os.path.abspath(Mask_RCNN_ROOT)

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import visualize


# In[ ]:


CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
               'bus', 'train', 'truck', 'boat', 'traffic light',
               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
               'kite', 'baseball bat', 'baseball glove', 'skateboard',
               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
               'teddy bear', 'hair drier', 'toothbrush']

VIP_CLASS = ['person','skis','snowboard']

# IMAGE_SHAPE = (467, 700, 3)
IMAGE_SHAPE = (667, 1000, 3)


# ---

# In[ ]:


def Show_Img(obj, showBox=True, showMask=True, getArray=False):
    """
    Show image for given image ID.

    Parameters (Input)
    ----------
    obj      : DataFrame, Series, str
        The Mask R-CNN record for a image
        or the path to the image file
    showBox  : bool
        Show the Boxes generated by Mask R-CNN
    showMask : bool
        Show the Masks generated by Mask R-CNN
    getArray : bool
        Return Array, not show image, will overwrite
        showBox=False, showMask=False

    Returns
    -------
    None (Just show the image)
    or
    Numpy array of the image
    """
    assert isinstance( obj, (str, pd.DataFrame, pd.Series) ), 'Input should be a Pandas DataFrame or Series.'

    if   isinstance(obj, str):
        imgFile = obj
    elif isinstance(obj['imgID'], str):
        imgFile = obj['imgID']
    elif isinstance(obj['imgID'], pd.Series):
        imgFile = obj['imgID'].unique()[0]
        obj = obj.where(obj['imgID']==imgFile).dropna()
    else:
        assert isinstance( obj['imgID'], (str, pd.Series) ), 'Unable to process:' + type(obj['imgID'])

    if not os.path.exists(imgFile):
        assert None, 'Not such image! ' + imgFile

    image  = skIR(imgFile)

    if getArray:
        return np.array(image)

    if isinstance(obj, str):
        return visualize.display_instances(
            image,
            np.zeros((2,2)), # Placeholder, rois
            np.zeros((2,2)), # Placeholder, masks
            np.zeros((2,2)), # Placeholder, class_ids
            np.array(0),     # Placeholder, CLASS_NAMES
            np.array(0),     # Placeholder, scores
            figsize=(8,8),
            show_mask=False,
            show_bbox=False,
        )
    else:
        result = {}
        if isinstance( obj, pd.DataFrame ):
            result['class_ids'] = np.array( obj['class_ids'].to_list() )
            result['scores']    = np.array( obj['scores'].to_list() )
            result['rois']      = np.array( obj[['x1','y1','x2','y2']].values)
        else:
            result['class_ids'] = np.array([obj['class_ids']])
            result['scores']    = np.array([obj['scores']])
            result['rois']      = np.array( obj[['x1','y1','x2','y2']].values)[np.newaxis, :]

        if showMask:
            result['masks'] = pd.Series(obj['masks']).apply( lambda row: list(map(int, list(row))) ).tolist()
            result['masks'] = np.rollaxis(
                    np.array(result['masks']).reshape(-1, IMAGE_SHAPE[0], IMAGE_SHAPE[1]), 0, 3
            ).astype(bool)
        else:
            result['masks'] = np.zeros((IMAGE_SHAPE[0], IMAGE_SHAPE[1], result['scores'].shape[0]))

        return visualize.display_instances(
            image,
            result['rois'],
            result['masks'],
            result['class_ids'].astype(int),
            CLASS_NAMES,
            result['scores'],
            figsize=(8,8),
            show_mask=showMask,
            show_bbox=showBox,
        )


# ---

# In[ ]:


def extInBoxPixels(obj, getMask=False, show=False):
    """
    Extract InBox pixels from given Box and image ID.

    Parameters (Input)
    ----------
    obj     : Series
        The record for a box
    getMask : bool
        Only extract the InMask pixels
    show    : bool
        Show the extracted pixels

    Returns
    -------
    ext_Box : Array
    Numpy array (Matrix) with InBox pixels
        Shape = (Unknown, Unknown, 3)
    """

    assert isinstance( obj, pd.Series ), 'Input should be a Pandas Series.'

    imgFile = obj['imgID']

    if not os.path.exists(imgFile):
        assert None, 'Not such image!'

    image = skIR(imgFile)

    (x1, y1, x2, y2) = obj[['x1','y1','x2','y2']].map(int)

    # Check image shape
    if image.shape != IMAGE_SHAPE:
        # Some are vertical image
        image = np.swapaxes(image,0,1)
        # Check again
        if image.shape != IMAGE_SHAPE:
            return None # Placehoder

    if not getMask:
        ext_Box = image[x1:x2, y1:y2, :]
    else:
        # Mask Invert
        ext_Mask = np.invert(
            np.array(
                pd.Series(obj['masks'])
                .apply( lambda row: list(map(int, list(row))) )
                .tolist()
            ).reshape(-1, IMAGE_SHAPE[0], IMAGE_SHAPE[1]).astype(bool)[0]
        )
        # First, Make Inverted Mask as a white/snow background (255,255,255)
        # Next,  Add image to the Inverted Mask
        # Then,  Clip the overflow (>255) pixels (make them white)
        ext_Img = (
            255*np.stack( [ext_Mask]*3, axis=2 )
            +image
        ).clip(max=255)
        # Finally, Crop the box
        ext_Box  = ext_Img[x1:x2, y1:y2, :]

    if show:
        plt.imshow(ext_Box)
    return ext_Box


# ---

# In[ ]:


def squareBox (BoxArray):
    """
    Reshape a Unknow shape Box with pixels to a square Box with 150x150.

    Parameters (Input)
    ----------
    BoxArray : numpy array (Matrix)
        Array with InBox pixels

    Returns
    -------
    BoxArraySquared : Array
    Numpy array (Matrix) with InBox pixels
        Shape = (150, 150, 3)
    """
    assert isinstance( BoxArray, np.ndarray ), 'Input should be a Numpy array.'


    BoxArraySquared = np.array(
        resize_tool(
            Image.fromarray(BoxArray.astype('uint8')),
            width  = 150,
            height = 150,
        )
    )
    return BoxArraySquared

################################################################################
def resize_tool(image_pil, width, height):
    '''
    Resize PIL image keeping ratio and using white background.

    From https://stackoverflow.com/questions/44370469/python-image-resizing-keep-proportion-add-white-background
    '''
    ratio_w = width  / image_pil.width
    ratio_h = height / image_pil.height
    if ratio_w < ratio_h:
        # It must be fixed by width
        resize_width = width
        resize_height = round(ratio_w * image_pil.height)
    else:
        # Fixed by height
        resize_width = round(ratio_h * image_pil.width)
        resize_height = height
    image_resize = image_pil.resize((resize_width, resize_height), Image.ANTIALIAS)

    background = Image.new('RGB', (width, height), (255, 255, 255))
    offset = (round((width - resize_width) / 2), round((height - resize_height) / 2))

    background.paste(image_resize, offset)

    return background


# In[ ]: