datagenerator.py

"""
Normalized Object Coordinate Space for Category-Level 6D Object Pose and Size Estimation
Jointly training for CAMERA, COCO, and REAL datasets 

Modified based on Mask R-CNN(https://github.com/matterport/Mask_RCNN)
Written by He Wang
"""

import os
import random
import math
import datetime
import re
import logging
from collections import OrderedDict
import numpy as np
import scipy.misc
import tensorflow as tf
import keras
import keras.backend as K
import keras.layers as KL
import keras.engine as KE
import keras.models as KM
import sys
#from keras.utils.np_utils import to_categorical

import utils
# import model
# from model import compose_image_meta

# Requires TensorFlow 1.3+ and Keras 2.0.8+.
from distutils.version import LooseVersion
assert LooseVersion(tf.__version__) >= LooseVersion("1.3")
assert LooseVersion(keras.__version__) >= LooseVersion('2.0.8')

def compose_image_meta(image_id, image_shape, window, active_class_ids):
	"""Takes attributes of an image and puts them in one 1D array. Use
	parse_image_meta() to parse the values back.
	
	image_id: An int ID of the image. Useful for debugging.
	image_shape: [height, width, channels]
	window: (y1, x1, y2, x2) in pixels. The area of the image where the real
			image is (excluding the padding)
	active_class_ids: List of class_ids available in the dataset from which
		the image came. Useful if training on images from multiple datasets
		where not all classes are present in all datasets.
	"""
	meta = np.array(
		[image_id] +            # size=1
		list(image_shape) +     # size=3
		list(window) +          # size=4 (y1, x1, y2, x2) in image cooredinates
		list(active_class_ids)  # size=num_classes
	)
	return meta

# Two functions (for Numpy and TF) to parse image_meta tensors.
def parse_image_meta(meta):
	"""Parses an image info Numpy array to its components.
	See compose_image_meta() for more details.
	"""
	image_id = meta[:, 0]
	image_shape = meta[:, 1:4]
	window = meta[:, 4:8]   # (y1, x1, y2, x2) window of image in in pixels
	active_class_ids = meta[:, 8:]
	return image_id, image_shape, window, active_class_ids


def parse_image_meta_graph(meta):
	"""Parses a tensor that contains image attributes to its components.
	See compose_image_meta() for more details.

	meta: [batch, meta length] where meta length depends on NUM_CLASSES
	"""
	image_id = meta[:, 0]
	image_shape = meta[:, 1:4]
	window = meta[:, 4:8]
	active_class_ids = meta[:, 8:]
	return [image_id, image_shape, window, active_class_ids]


def mold_image(images, config):
	"""
	Takes RGB images with 0-255 values and subtraces
	the mean pixel and converts it to float. Expects image
	colors in RGB order.
	"""
	return images.astype(np.float32) - config.MEAN_PIXEL


def unmold_image(normalized_images, config):
	"""Takes a image normalized with mold() and returns the original."""
	return (normalized_images + config.MEAN_PIXEL).astype(np.uint8)

############################################################
#  Data Generator
############################################################

def load_image_gt(dataset, config, image_id, augment=False,
                  use_mini_mask=False, load_scale=False):
    """Load and return ground truth data for an image (image, mask, bounding boxes).

    augment: If true, apply random image augmentation. Currently, only
        horizontal flipping is offered.
    use_mini_mask: If False, returns full-size masks that are the same height
        and width as the original image. These can be big, for example
        1024x1024x100 (for 100 instances). Mini masks are smaller, typically,
        224x224 and are generated by extracting the bounding box of the
        object and resizing it to MINI_MASK_SHAPE.

    Returns:
    image: [height, width, 3]
    shape: the original shape of the image before resizing and cropping.
    bbox: [instance_count, (y1, x1, y2, x2, class_id)]
    mask: [height, width, instance_count]. The height and width are those
        of the image unless use_mini_mask is True, in which case they are
        defined in MINI_MASK_SHAPE.
    """
    # Load image and mask

    if config.TRAINING_AUGMENTATION and dataset.subset == 'train':
        # image, mask, coord, normal, class_ids, scales, domain_label = dataset.load_augment_data(image_id)
        # image, mask, coord, class_ids, domain_label = dataset.load_augment_data(image_id) #scales
        image, mask, coord, class_ids, scales, domain_label = dataset.load_augment_data(image_id) #scales
    else:
        image = dataset.load_image(image_id)
        # mask, coord, normal, class_ids, scales, domain_label = dataset.load_mask(image_id)
        mask, coord, class_ids, domain_label = dataset.load_mask(image_id) #scales

    #print('maximum mask: ', np.amax(mask))
    #print('maximum coord: ', np.amax(coord))
    shape = image.shape
    image, window, scale, padding = utils.resize_image(
        image, 
        min_dim=config.IMAGE_MIN_DIM, 
        max_dim=config.IMAGE_MAX_DIM,
        padding=config.IMAGE_PADDING)
    mask  = utils.resize_mask(mask, scale, padding)
    coord = utils.resize_mask(coord, scale, padding)
    
    #print('maximum mask after resize: ', np.amax(mask))
    #print('maximum coord after resize: ', np.amax(coord))

    # Random horizontal flips.
    #if augment:
    #    if random.randint(0, 1):
    #        image = np.fliplr(image)
    #        mask = np.fliplr(mask)
    #        coord = np.fliplr(coord)

    # Bounding boxes. Note that some boxes might be all zeros
    # if the corresponding mask got cropped out.
    # bbox: [num_instances, (y1, x1, y2, x2)]
    bbox = utils.extract_bboxes(mask)

    # Add class_id as the last value in bbox
    bbox = np.hstack((bbox, class_ids[:, np.newaxis]))
    #print(class_ids)

    # Active classes
    # Different datasets have different classes, so track the
    # classes supported in the dataset of this image.
    #active_class_ids = np.zeros([dataset.num_classes], dtype=np.int32)
    #class_ids_2 = dataset.source_class_ids[dataset.image_info[image_id]["source"]]
    #print(class_ids_2)

    ####assert set(class_ids) == set(class_ids_2)
    #active_class_ids[class_ids_2] = 1
    active_class_ids = np.ones([dataset.num_classes], dtype=np.int32)

    # Resize masks to smaller size to reduce memory usage
    if use_mini_mask:
        #print('\n')
        #print(bbox)
        mask  =  utils.minimize_mask(bbox, mask, config.MINI_MASK_SHAPE)
        coord =  utils.minimize_mask(bbox, coord, config.MINI_MASK_SHAPE)
        # normal =  utils.minimize_mask(bbox, normal, config.MINI_MASK_SHAPE)
        #print('maximum mask after mini mask: ', np.amax(mask))
        #print('maximum coord after mini mask: ', np.amax(coord))

    # Image meta data
    image_meta = compose_image_meta(image_id, shape, window, active_class_ids)


    if load_scale:
        # return image, image_meta, bbox, mask, coord, normal, domain_label, scales
        return image, image_meta, bbox, mask, coord, domain_label, scales
    else:
        # return image, image_meta, bbox, mask, coord, normal, domain_label
        return image, image_meta, bbox, mask, coord, domain_label


def build_detection_targets(rpn_rois, gt_boxes, gt_masks, config):
    """Generate targets for training Stage 2 classifier and mask heads.

    Inputs:
    rpn_rois: [N, (y1, x1, y2, x2)] proposal boxes.
    gt_boxes: [instance count, (y1, x1, y2, x2, class_id)]
    gt_masks: [height, width, instance count] Grund truth masks. Can be full
              size or mini-masks.

    Returns:
    rois: [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)]
    class_ids: [TRAIN_ROIS_PER_IMAGE]. Int class IDs.
    bboxes: [TRAIN_ROIS_PER_IMAGE, NUM_CLASSES, 5]. Rows are class-specific 
            bbox refinments [y, x, log(h), log(w), weight].
    masks: [TRAIN_ROIS_PER_IMAGE, height, width, NUM_CLASSES). Class specific masks cropped
           to bbox boundaries and resized to neural network output size.
    """
    assert rpn_rois.shape[0] > 0
    assert gt_boxes.dtype == np.int32, "Expected int but got {}".format(gt_boxes.dtype)
    assert gt_masks.dtype == np.bool_, "Expected bool but got {}".format(gt_masks.dtype)

    # It's common to add GT Boxes to ROIs but we don't do that here because
    # according to XinLei Chen's paper, it doesn't help.

    # Trim empty padding in gt_boxes and gt_masks parts
    instance_ids = np.where(gt_boxes[:, 4] > 0)[0]
    assert instance_ids.shape[0] > 0, "Image must contain instances."
    gt_boxes = gt_boxes[instance_ids]
    gt_masks = gt_masks[:, :, instance_ids]

    # Compute areas of ROIs and ground truth boxes.
    rpn_roi_area = (rpn_rois[:, 2] - rpn_rois[:, 0]) * (rpn_rois[:, 3] - rpn_rois[:, 1])
    gt_box_area = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] - gt_boxes[:, 1])

    # Compute overlaps [rpn_rois, gt_boxes]
    overlaps = np.zeros((rpn_rois.shape[0], gt_boxes.shape[0]))
    for i in range(overlaps.shape[1]):
        gt = gt_boxes[i][:4]
        overlaps[:,i] = utils.compute_iou(gt, rpn_rois, gt_box_area[i], rpn_roi_area)

    # Assign ROIs to GT boxes
    rpn_roi_iou_argmax = np.argmax(overlaps, axis=1)
    rpn_roi_iou_max = overlaps[np.arange(overlaps.shape[0]), rpn_roi_iou_argmax]
    rpn_roi_gt_boxes = gt_boxes[rpn_roi_iou_argmax]  # GT box assigned to each ROI

    # Positive ROIs are those with >= 0.5 IoU with a GT box. 
    fg_ids = np.where(rpn_roi_iou_max > 0.5)[0]

    # Negative ROIs are those with max IoU 0.1-0.5 (hard example mining)
    # TODO: To hard example mine or not to hard example mine, that's the question
#     bg_ids = np.where((rpn_roi_iou_max >= 0.1) & (rpn_roi_iou_max < 0.5))[0]
    bg_ids = np.where(rpn_roi_iou_max < 0.5)[0]

    # Subsample ROIs. Aim for 33% foreground.
    # FG
    fg_roi_count = int(config.TRAIN_ROIS_PER_IMAGE * config.ROI_POSITIVE_RATIO)
    if fg_ids.shape[0] > fg_roi_count:
        keep_fg_ids = np.random.choice(fg_ids, fg_roi_count, replace=False)
    else:
        keep_fg_ids = fg_ids
    # BG
    remaining = config.TRAIN_ROIS_PER_IMAGE - keep_fg_ids.shape[0]
    if bg_ids.shape[0] > remaining:
        keep_bg_ids = np.random.choice(bg_ids, remaining, replace=False)
    else:
        keep_bg_ids = bg_ids
    # Combine indicies of ROIs to keep
    keep = np.concatenate([keep_fg_ids, keep_bg_ids])
    # Need more?
    remaining = config.TRAIN_ROIS_PER_IMAGE - keep.shape[0]
    if remaining > 0:
        # Looks like we don't have enough samples to maintain the desired
        # balance. Reduce requirements and fill in the rest. This is
        # likely different from the Mask RCNN paper.

        # There is a small chance we have neither fg nor bg samples.
        if keep.shape[0] == 0:
            # Pick bg regions with easier IoU threshold
            bg_ids = np.where(rpn_roi_iou_max < 0.5)[0]
            assert bg_ids.shape[0] >= remaining
            keep_bg_ids = np.random.choice(bg_ids, remaining, replace=False)
            assert keep_bg_ids.shape[0] == remaining
            keep = np.concatenate([keep, keep_bg_ids])
        else:
            # Fill the rest with repeated bg rois.
            keep_extra_ids = np.random.choice(keep_bg_ids, remaining, replace=True)
            keep = np.concatenate([keep, keep_extra_ids])
    assert keep.shape[0] == config.TRAIN_ROIS_PER_IMAGE, \
        "keep doesn't match ROI batch size {}, {}".format(
            keep.shape[0], config.TRAIN_ROIS_PER_IMAGE)

    # Reset the gt boxes assigned to BG ROIs.
    rpn_roi_gt_boxes[keep_bg_ids, :] = 0

    # For each kept ROI, assign a class_id, and for FG ROIs also add bbox refinement.
    rois = rpn_rois[keep, :4]
    roi_gt_boxes = rpn_roi_gt_boxes[keep]
    class_ids = roi_gt_boxes[:,4].astype(np.int32)
    roi_gt_assignment = rpn_roi_iou_argmax[keep]

    # Class-aware bbox shifts. [y, x, log(h), log(w), weight]. Weight is 0 or 1 to
    # determine if a bbox is included in the loss.
    bboxes = np.zeros((config.TRAIN_ROIS_PER_IMAGE, config.NUM_CLASSES, 5), dtype=np.float32)
    pos_ids = np.where(class_ids > 0)[0]
    bboxes[pos_ids, class_ids[pos_ids], :4] = utils.box_refinement(rois[pos_ids], roi_gt_boxes[pos_ids, :4])
    bboxes[pos_ids, class_ids[pos_ids], 4] = 1  # weight = 1 to influence the loss
    # Normalize bbox refinments
    bboxes[:, :, :4] /= config.BBOX_STD_DEV

    # Generate class-specific target masks.
    masks = np.zeros((config.TRAIN_ROIS_PER_IMAGE, config.MASK_SHAPE[0], config.MASK_SHAPE[1], config.NUM_CLASSES), 
                     dtype=np.float32)
    for i in pos_ids:
        class_id = class_ids[i]
        assert class_id > 0, "class id must be greater than 0"
        gt_id = roi_gt_assignment[i]
        class_mask = gt_masks[:, :, gt_id]
        
        if config.USE_MINI_MASK:
            # Create a mask placeholder, the size of the image
            placeholder = np.zeros(config.IMAGE_SHAPE[:2], dtype=bool)
            # GT box
            gt_y1, gt_x1, gt_y2, gt_x2 = gt_boxes[gt_id][:4]
            gt_w = gt_x2 - gt_x1
            gt_h = gt_y2 - gt_y1
            # Resize mini mask to size of GT box
            placeholder[gt_y1:gt_y2, gt_x1:gt_x2] = \
                np.round(scipy.misc.imresize(class_mask.astype(float), (gt_h, gt_w), 
                                             interp='nearest') / 255.0).astype(bool)
            # Place the mini batch in the placeholder
            class_mask = placeholder
            
        # Pick part of the mask and resize it
        y1, x1, y2, x2 = rois[i][:4].astype(np.int32)
        m = class_mask[y1:y2, x1:x2]
        mask = scipy.misc.imresize(m.astype(float), config.MASK_SHAPE, interp='nearest') / 255.0
        masks[i,:,:,class_id] = mask
        
    return rois, class_ids, bboxes, masks


def build_rpn_targets(image_shape, anchors, gt_boxes, config):
    """Given the anchors and GT boxes, compute overlaps and identify positive
    anchors and deltas to refine them to match their corresponding GT boxes.

    anchors: [num_anchors, (y1, x1, y2, x2)]
    gt_boxes: [num_gt_boxes, (y1, x1, y2, x2, class_id)]

    Returns:
    rpn_match: [N] (int32) matches between anchors and GT boxes.
               1 = positive anchor, -1 = negative anchor, 0 = neutral
    rpn_bbox: [N, (dy, dx, log(dh), log(dw))] Anchor bbox deltas.
    """
    # RPN Match: 1 = positive anchor, -1 = negative anchor, 0 = neutral
    rpn_match = np.zeros([anchors.shape[0]], dtype=np.int32)
    # RPN bounding boxes: [max anchors per image, (dy, dx, log(dh), log(dw))]
    rpn_bbox = np.zeros((config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4))

    # Areas of anchors and GT boxes
    gt_box_area = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] - gt_boxes[:, 1])
    anchor_area = (anchors[:, 2] - anchors[:, 0]) * (anchors[:, 3] - anchors[:, 1])

    # Compute overlaps [num_anchors, num_gt_boxes]
    # Each cell contains the IoU of an anchor and GT box.
    overlaps = np.zeros((anchors.shape[0], gt_boxes.shape[0]))
    for i in range(overlaps.shape[1]):
        gt = gt_boxes[i][:4]
        overlaps[:,i] = utils.compute_iou(gt, anchors, gt_box_area[i], anchor_area)

    # Match anchors to GT Boxes
    # If an anchor overlaps a GT box with IoU >= 0.7 then it's positive.
    # If an anchor overlaps a GT box with IoU < 0.3 then it's negative.
    # Neutral anchors are those that don't match the conditions above, 
    # and they don't influence the loss function.
    # However, don't keep any GT box unmatched (rare, but happens). Instead,
    # match it to the closest anchor (even if its max IoU is < 0.3).
    #
    # 1. Set negative anchors first. It gets overwritten if a gt box is matched to them.
    anchor_iou_argmax = np.argmax(overlaps, axis=1)
    anchor_iou_max = overlaps[np.arange(overlaps.shape[0]), anchor_iou_argmax]
    rpn_match[anchor_iou_max < 0.3] = -1
    # 2. Set an anchor for each GT box (regardless of IoU value).
    # TODO: If multiple anchors have the same IoU match all of them
    gt_iou_argmax = np.argmax(overlaps, axis=0)
    rpn_match[gt_iou_argmax] = 1
    # 3. Set anchors with high overlap as positive.
    rpn_match[anchor_iou_max >= 0.7] = 1

    # Subsample to balance positive and negative anchors
    # Don't let positives be more than half the anchors
    ids = np.where(rpn_match == 1)[0]
    extra = len(ids) - (config.RPN_TRAIN_ANCHORS_PER_IMAGE // 2)
    if extra > 0:
        # Reset the extra ones to neutral
        ids = np.random.choice(ids, extra, replace=False)
        rpn_match[ids] = 0
    # Same for negative proposals
    ids = np.where(rpn_match == -1)[0]
    extra = len(ids) - (config.RPN_TRAIN_ANCHORS_PER_IMAGE - np.sum(rpn_match == 1))
    if extra > 0:
        # Rest the extra ones to neutral
        ids = np.random.choice(ids, extra, replace=False)
        rpn_match[ids] = 0

    # For positive anchors, compute shift and scale needed to transform them
    # to match the corresponding GT boxes.
    ids = np.where(rpn_match == 1)[0]
    ix = 0  # index into rpn_bbox
    # TODO: use box_refinment() rather that duplicating the code here
    for i, a in zip(ids, anchors[ids]):
        # Closest gt box (it might have IoU < 0.7)
        gt = gt_boxes[anchor_iou_argmax[i], :4]

        # Convert coordinates to center plus width/height.
        # GT Box
        gt_h = gt[2] - gt[0]
        gt_w = gt[3] - gt[1]
        gt_center_y = gt[0] + 0.5 * gt_h
        gt_center_x = gt[1] + 0.5 * gt_w
        # Anchor
        a_h = a[2] - a[0]
        a_w = a[3] - a[1]
        a_center_y = a[0] + 0.5 * a_h
        a_center_x = a[1] + 0.5 * a_w

        # Compute the bbox refinement that the RPN should predict.
        rpn_bbox[ix] = [
            (gt_center_y - a_center_y) / a_h,
            (gt_center_x - a_center_x) / a_w,
            np.log(gt_h / a_h),
            np.log(gt_w / a_w),
        ]
        # Normalize
        rpn_bbox[ix] /= config.RPN_BBOX_STD_DEV
        ix += 1

    return rpn_match, rpn_bbox


def generate_random_rois(image_shape, count, gt_boxes):
    """Generates ROI proposals similar to what a region proposal network
    would generate.

    image_shape: [Height, Width, Depth]
    count: Number of ROIs to generate
    gt_boxes: [N, (y1, x1, y2, x2, class_id)] Ground trugh boxes in pixels.

    Returns: [count, (y1, x1, y2, x2)] ROI boxes in pixels.
    """
    # placeholder
    rois = np.zeros((count, 4), dtype=np.int32)
    
    # Generate random ROIs around GT boxes (90% of count)
    rois_per_box = int(0.9 * count / gt_boxes.shape[0])
    for i in range(gt_boxes.shape[0]):
        gt_y1, gt_x1, gt_y2, gt_x2 = gt_boxes[i,:4]
        h = gt_y2 - gt_y1
        w = gt_x2 - gt_x1
        # random boundaries
        r_y1 = max(gt_y1-h, 0)
        r_y2 = min(gt_y2+h, image_shape[0])
        r_x1 = max(gt_x1-w, 0)
        r_x2 = min(gt_x2+w, image_shape[1])
        
        # To avoid generating boxes with zero area, we generate double what
        # we need and filter out the extra. If we get fewer valid boxes 
        # than we need, we loop and try again.
        while True:
            y1y2 = np.random.randint(r_y1, r_y2, (rois_per_box*2, 2))
            x1x2 = np.random.randint(r_x1, r_x2, (rois_per_box*2, 2))
            # Filter out zero area boxes
            threshold = 1
            y1y2 = y1y2[np.abs(y1y2[:,0] - y1y2[:,1]) >= threshold][:rois_per_box]
            x1x2 = x1x2[np.abs(x1x2[:,0] - x1x2[:,1]) >= threshold][:rois_per_box]
            if y1y2.shape[0] == rois_per_box and x1x2.shape[0] == rois_per_box:
                break
        
        # Sort on axis 1 to ensure x1 <= x2 and y1 <= y2 and then reshape
        # into x1, y1, x2, y2 order
        x1, x2 = np.split(np.sort(x1x2, axis=1), 2, axis=1)
        y1, y2 = np.split(np.sort(y1y2, axis=1), 2, axis=1)
        box_rois = np.hstack([y1, x1, y2, x2])
        rois[rois_per_box*i:rois_per_box*(i+1)] = box_rois
    
    # Generate random ROIs anywhere in the image (10% of count)
    remaining_count = count - (rois_per_box * gt_boxes.shape[0])
    # To avoid generating boxes with zero area, we generate double what
    # we need and filter out the extra. If we get fewer valid boxes 
    # than we need, we loop and try again.
    while True:
        y1y2 = np.random.randint(0, image_shape[0], (remaining_count * 2, 2))
        x1x2 = np.random.randint(0, image_shape[1], (remaining_count * 2, 2))
        # Filter out zero area boxes
        threshold = 1
        y1y2 = y1y2[np.abs(y1y2[:,0] - y1y2[:,1]) >= threshold][:remaining_count]
        x1x2 = x1x2[np.abs(x1x2[:,0] - x1x2[:,1]) >= threshold][:remaining_count]
        if y1y2.shape[0] == remaining_count and x1x2.shape[0] == remaining_count:
            break
    
    # Sort on axis 1 to ensure x1 <= x2 and y1 <= y2 and then reshape
    # into x1, y1, x2, y2 order
    x1, x2 = np.split(np.sort(x1x2, axis=1), 2, axis=1)
    y1, y2 = np.split(np.sort(y1y2, axis=1), 2, axis=1)
    global_rois = np.hstack([y1, x1, y2, x2])
    rois[-remaining_count:] = global_rois
    return rois


def data_generator(dataset, config, shuffle=True, augment=False, random_rois=0,
                   batch_size=1, detection_targets=False):
    """A generator that returns images and corresponding target class ids, 
    bounding box deltas, and masks.

    dataset: The Dataset object to pick data from
    config: The model config object
    shuffle: If True, shuffles the samples before every epoch
    augment: If True, applies image augmentation to images (currently only 
             horizontal flips are supported)
    random_rois: If > 0 then generate proposals to be used to train the
                 network classifier and mask heads. Useful if training
                 the Mask RCNN part without the RPN.
    batch_size: How many images to return in each call
    detection_targets: If True, generate detection targets (class IDs, bbox
        deltas, and masks). Typically for debugging or visualizations because
        in trainig detection targets are generated by DetectionTargetLayer.

    Returns a Python generator. Upon calling next() on it, the 
    generator returns two lists, inputs and outputs. The containtes
    of the lists differs depending on the received arguments:
    inputs list:
    - images: [batch, H, W, C]
    - image_meta: [batch, size of image meta]
    - rpn_match: [batch, N] Integer (1=positive anchor, -1=negative, 0=neutral)
    - rpn_bbox: [batch, N, (dy, dx, log(dh), log(dw))] Anchor bbox deltas.
    - gt_boxes: [batch, MAX_GT_INSTANCES, (y1, x1, y2, x2, class_id)]
    - gt_masks: [batch, height, width, MAX_GT_INSTANCES]. The height and width
                are those of the image unless use_mini_mask is True, in which
                case they are defined in MINI_MASK_SHAPE.
    - gt_coords: [batch, height, width, MAX_GT_INSTANCES, 3]. The height and width
                are those of the image unless use_mini_mask is True, in which
                case they are defined in MINI_MASK_SHAPE.
    - gt_normals: [batch, height, width, MAX_GT_INSTANCES, 3]. The height and width
                are those of the image unless use_mini_mask is True, in which
                case they are defined in MINI_MASK_SHAPE.

    outputs list: Usually empty in regular training. But if detection_targets
        is True then the outputs list contains target class_ids, bbox deltas,
        and masks.
    """


    b = 0  # batch item index
    #image_index = -1
    #image_ids = np.copy(dataset.image_ids)


    source_image_ids = dataset.source_image_ids.copy()
    sources = source_image_ids.keys()
    print(sources)
    
    # ShapeNetTOI, Real, coco
    ## Determine how to choose betweem dataset 
    #source_names = ['CAMERA', 'Real', 'coco']
    source_names = ['SOM', 'coco', 'open_images']

    weight = np.array([0, 0, 0], dtype=np.float32) 
    weight_consts = config.SOURCE_WEIGHT #[3, 1, 1]
    weight_sum = 0.0

    for i, source_name in enumerate(source_names):
        if source_name in sources:
            weight[i] = weight_consts[i]
            weight_sum += weight_consts[i]

    weight = weight/weight_sum
    assert np.sum(weight) == 1, "[ Error ]: total sum of weights is {} != 1".format(np.sum(weight))

    source_image_index = {}
    for source_name in sources:
        source_image_index[source_name] = -1


    error_count = 0

    # Anchors
    # [anchor_count, (y1, x1, y2, x2)]
    anchors = utils.generate_pyramid_anchors(config.RPN_ANCHOR_SCALES,
                                                 config.RPN_ANCHOR_RATIOS,
                                                 config.BACKBONE_SHAPES,
                                                 config.BACKBONE_STRIDES,
                                                 config.RPN_ANCHOR_STRIDE)


    # Keras requires a generator to run indefinately.
    while True:
        try:
            # draw a random number from 0, 1, 2 according to weight
            source_ind = np.random.choice([0, 1, 2], 1, p=weight)[0]
            draw_source_name = source_names[source_ind]
            #print('draw_source_name: ', draw_source_name)

            # the image_ids list for the drawn source 
            draw_source_ids = source_image_ids[draw_source_name]

            # Increment index to pick next image.
            source_image_index[draw_source_name] = (source_image_index[draw_source_name] + 1) % len(draw_source_ids)
            image_id_in_draw_source = source_image_index[draw_source_name]

            ## shuffle if at the start of an epoch.
            if shuffle and image_id_in_draw_source == 0:
                np.random.shuffle(source_image_ids[draw_source_name])

            # image_index = (image_index + 1) % len(image_ids)
            # ## shuffle if at the start of an epoch.
            # if shuffle and image_index == 0:
            #     np.random.shuffle(image_ids)

            # Get GT bounding boxes and masks for image.
            # image_id = image_ids[image_index]

            image_id = source_image_ids[draw_source_name][image_id_in_draw_source]
            #print('image_id_in_draw_source: ', image_id_in_draw_source)
            #print('image_id: ', image_id)

            # image, image_meta, gt_boxes, gt_masks, gt_coords, gt_normals, gt_domain_label = \
            #     load_image_gt(dataset, config, image_id, augment=augment, use_mini_mask=config.USE_MINI_MASK)
            image, image_meta, gt_boxes, gt_masks, gt_coords, gt_domain_label = \
                load_image_gt(dataset, config, image_id, augment=augment, use_mini_mask=config.USE_MINI_MASK)
            
            # ones-NOTE: initialse ones. default dtype is float64
            # input_ones = np.ones((config.TRAIN_ROIS_PER_IMAGE*config.NUM_CLASSES, 
            #                       config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1], config.K*config.K, 1))

            # Skip images that have no instances. This can happen in cases 
            # where we train on a subset of classes and the image doesn't
            # have any of the classes we care about.
            if np.sum(gt_boxes) <= 0:
                continue

            # RPN Targets
            rpn_match, rpn_bbox = build_rpn_targets(image.shape, anchors, gt_boxes, config)

            # Mask R-CNN Targets
            if random_rois:
                rpn_rois = generate_random_rois(image.shape, random_rois, gt_boxes)
                # TODO: mrcnn_coord for detection_targets = True
                if detection_targets:
                    # Append two columns of zeros. TODO: needed?
                    rpn_rois = np.hstack([rpn_rois, np.zeros([rpn_rois.shape[0], 2], dtype=np.int32)])
                    rois, mrcnn_class_ids, mrcnn_bbox, mrcnn_mask =\
                        build_detection_targets(rpn_rois, gt_boxes, gt_masks, config)

            # Init batch arrays
            if b == 0:
                batch_image_meta = np.zeros((batch_size,)+image_meta.shape, dtype=image_meta.dtype)
                batch_rpn_match = np.zeros([batch_size, anchors.shape[0], 1], dtype=rpn_match.dtype)
                batch_rpn_bbox = np.zeros([batch_size, config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4], dtype=rpn_bbox.dtype)
                batch_images = np.zeros((batch_size,)+image.shape, dtype=np.float32)
                batch_gt_boxes = np.zeros((batch_size, config.MAX_GT_INSTANCES, 5), dtype=np.int32)
                batch_gt_domain_labels = np.zeros((batch_size, 1), dtype=np.bool_)

                


                if config.USE_MINI_MASK:
                    batch_gt_masks = np.zeros((batch_size, config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1], 
                                               config.MAX_GT_INSTANCES))
                    batch_gt_coords = np.zeros((batch_size, config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1],
                                               config.MAX_GT_INSTANCES, 3), dtype=np.float32)
                    # batch_gt_normals = np.zeros((batch_size, config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1],
                    #                                             config.MAX_GT_INSTANCES, 3), dtype=np.float32)                                                     

                    # Ones-NOTE: add which dimension to put GT instance ?
                    # batch_ones = np.zeros((batch_size, config.TRAIN_ROIS_PER_IMAGE*config.NUM_CLASSES, 
                    #                        config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1], config.K*config.K, 1))

                else:
                    batch_gt_masks = np.zeros((batch_size, image.shape[0], image.shape[1], config.MAX_GT_INSTANCES))
                    batch_gt_coords = np.zeros((batch_size, image.shape[0], image.shape[1], config.MAX_GT_INSTANCES, 3))
                    
                    # batch_gt_normals = np.zeros((batch_size, image.shape[0], image.shape[1], config.MAX_GT_INSTANCES, 3))
                    # batch_ones = np.zeros( (batch_size, config.TRAIN_ROIS_PER_IMAGE*config.NUM_CLASSES, image.shape[0], image.shape[1],
                                            # config.K*config.K, 1))

                if random_rois:
                    batch_rpn_rois = np.zeros((batch_size,rpn_rois.shape[0], 4), dtype=rpn_rois.dtype)
                    if detection_targets:
                        batch_rois = np.zeros((batch_size,)+rois.shape, dtype=rois.dtype)
                        batch_mrcnn_class_ids = np.zeros((batch_size,)+mrcnn_class_ids.shape, dtype=mrcnn_class_ids.dtype)
                        batch_mrcnn_bbox = np.zeros((batch_size,)+mrcnn_bbox.shape, dtype=mrcnn_bbox.dtype)
                        batch_mrcnn_mask = np.zeros((batch_size,)+mrcnn_mask.shape, dtype=mrcnn_mask.dtype)

            # If more instances than fits in the array, sub-sample from them.
            if gt_boxes.shape[0] > config.MAX_GT_INSTANCES:
                ids = np.random.choice(np.arange(gt_boxes.shape[0]), config.MAX_GT_INSTANCES, replace=False)
                gt_boxes = gt_boxes[ids]
                gt_masks = gt_masks[:,:,ids]
                gt_coords = gt_coords[:,:,ids,:]
                # gt_normals = gt_normals[:,:,ids,:]

            # Add to batch
            batch_image_meta[b] = image_meta
            batch_rpn_match[b] = rpn_match[:, np.newaxis]
            batch_rpn_bbox[b] = rpn_bbox
            batch_images[b] = mold_image(image.astype(np.float32), config)
            batch_gt_boxes[b,:gt_boxes.shape[0]] = gt_boxes
            batch_gt_masks[b,:,:,:gt_masks.shape[-1]] = gt_masks
            batch_gt_coords[b, :, :, :gt_coords.shape[-2], :] = gt_coords
            # batch_gt_normals[b, :, :, :gt_normals.shape[-2], :] = gt_normals
            batch_gt_domain_labels[b] = gt_domain_label
            # batch_ones[b] = input_ones

            if random_rois:
                batch_rpn_rois[b] = rpn_rois[:,:4]
                if detection_targets:
                    batch_rois[b] = rois
                    batch_mrcnn_class_ids[b] = mrcnn_class_ids
                    batch_mrcnn_bbox[b] = mrcnn_bbox
                    batch_mrcnn_mask[b] = mrcnn_mask
            b += 1

            # Batch full?
            if b >= batch_size:
                # inputs = [batch_images, batch_image_meta, batch_ones, batch_rpn_match, batch_rpn_bbox,
                #           batch_gt_boxes, batch_gt_masks, batch_gt_coords, batch_gt_normals, 
                #           batch_gt_domain_labels]
                inputs = [batch_images, batch_image_meta, batch_rpn_match, batch_rpn_bbox,
                          batch_gt_boxes, batch_gt_masks, batch_gt_coords, batch_gt_domain_labels]

                outputs = []

                if random_rois:
                    inputs.extend([batch_rpn_rois])
                    if detection_targets:
                        inputs.extend([batch_rois])
                        # Keras requires that output and targets have the same number of dimensions
                        batch_mrcnn_class_ids = np.expand_dims(batch_mrcnn_class_ids, -1)
                        outputs.extend([batch_mrcnn_class_ids, batch_mrcnn_bbox, batch_mrcnn_mask])

                yield inputs, outputs

                # start a new batch
                b = 0
        except (GeneratorExit, KeyboardInterrupt):
            raise
        except:
            # Log it and skip the image
            logging.exception("Error processing image {}".format(dataset.image_info[draw_source_ids]))
            error_count += 1
            if error_count > 5:
                raise