data_feed.py

import os
import time
from collections import OrderedDict

import cv2
import numpy as np

__all__ = ['reader']


def preprocess(org_im, shrink):
    image = org_im.copy()
    image_height, image_width, image_channel = image.shape
    if shrink != 1:
        image_height, image_width = int(image_height * shrink), int(image_width * shrink)
        image = cv2.resize(image, (image_width, image_height), cv2.INTER_NEAREST)
    # HWC to CHW
    if len(image.shape) == 3:
        image = np.swapaxes(image, 1, 2)
        image = np.swapaxes(image, 1, 0)
    # mean, std
    mean = [104., 117., 123.]
    scale = 0.007843
    image = image.astype('float32')
    image -= np.array(mean)[:, np.newaxis, np.newaxis].astype('float32')
    image = image * scale
    return image, image_height, image_width


def reader(images, paths, shrink):
    """
    Preprocess to yield image.

    Args:
        images (list(numpy.ndarray)): images data, shape of each is [H, W, C], color space is BGR.
        paths (list[str]): paths to images.
        shrink (float): parameter to control the resize scale in preprocess.

    Yield:
        each (collections.OrderedDict): info of original image, preprocessed image.
    """
    component = list()
    if paths is not None:
        assert type(paths) is list, "paths should be a list."
        for im_path in paths:
            each = OrderedDict()
            assert os.path.isfile(im_path), "The {} isn't a valid file path.".format(im_path)
            im = cv2.imread(im_path)
            each['org_im'] = im
            each['org_im_path'] = im_path
            component.append(each)
    if images is not None:
        assert type(images) is list, "images should be a list."
        for im in images:
            each = OrderedDict()
            each['org_im'] = im
            each['org_im_path'] = 'ndarray_time={}'.format(round(time.time(), 6) * 1e6)
            component.append(each)

    for element in component:
        element['image'], element['image_height'], element['image_width'] = preprocess(element['org_im'], shrink)
        yield element