import numpy as np

def random_augmentation(img, mask):
    #you can add any augmentations you need
    return img, mask

def batch_generator(image, mask, 
                    batch_size=1, 
                    crop_size=0, 
                    patch_size=256,
                    bbox= None,
                    augmentation=False):
    '''
    image: nparray, must have 3 dimension
    mask: nparray, 2 dimensions, same size as image
    batch_size: int, number of images in a batch
    patch_size: int, size of the image returned, patch is square
    crop_size: int, how much pixels should be cropped off the mask
    bbox: None or tuple of 4 ints, (min_y, max_y, min_x, max_x), the data is selected from within the bbox
    augmentation: turn on/off data augmentation. The augmentation function is random_augmentation() above
    
    returns batch of image and mask patches, image is turned to 'channels last' as required by unet
    '''
    if np.ndim(mask) != 2 or np.ndim(image) != 3:
        raise ValueError('image must have 3 dims and mask 2 dims')
    if mask.shape != image.shape[1:]:
        raise ValueError('image and mask shape is different')
 
    im_max = float(np.max(image))
    mask_max = 1.0
    
    #select subimage
    if bbox is not None:
        # check bbox
        if bbox[0] < 0 or bbox [2] < 0 \
            or bbox[1] > mask.shape[0] or bbox[3] > mask.shape[0] \
            or bbox[0] + patch_size > bbox[1] or bbox[2] + patch_size > bbox[3] \
            or patch_size <= 0:
            raise ValueError("Incorrect bbox or patch size")
        img_ = image[:, bbox[0] : bbox[1], bbox[2]:bbox[3]]
        mask_ = mask[bbox[0] : bbox[1], bbox[2]:bbox[3]]
    else:
        img_ = image
        mask_ = mask
    while 1:
        x = []
        y = []
        for i in range (batch_size):
            random_x = np.random.randint(0, mask_.shape[1] - patch_size)
            random_y = np.random.randint(0, mask_.shape[0] - patch_size)
            img_patch = img_[:,
                  random_y : random_y + patch_size, 
                  random_x : random_x + patch_size] / im_max
            
            # transform the image from channels-first (rasterio format) to channels-last (default tensorflow format)
            img_patch = np.moveaxis(img_patch, 0, 2)
            
            mask_patch = mask_[random_y : random_y + patch_size, 
                  random_x : random_x + patch_size] / mask_max
            
            if augmentation:
                img_patch, mask_patch = random_augmentation(img_patch, mask_patch)
                
            # mask is cropped as it may be useful for some convnets that have output size less than input
            if crop_size > 0:
                mask_patch = mask_patch[crop_size : -crop_size, 
                                        crop_size : -crop_size]
            
            mask_patch = np.expand_dims(mask_patch, 2)
            x.append(img_patch)
            y.append(mask_patch)
        yield (np.array(x), np.array(y))
        

[((29, 19, 29, 32), 'numpy.max', 'np.max', ({(29, 26, 29, 31): 'image'}, {}), '(image)', True, 'import numpy as np\n'), ((24, 7, 24, 20), 'numpy.ndim', 'np.ndim', ({(24, 15, 24, 19): 'mask'}, {}), '(mask)', True, 'import numpy as np\n'), ((24, 29, 24, 43), 'numpy.ndim', 'np.ndim', ({(24, 37, 24, 42): 'image'}, {}), '(image)', True, 'import numpy as np\n'), ((49, 23, 49, 72), 'numpy.random.randint', 'np.random.randint', ({(49, 41, 49, 42): '0', (49, 44, 49, 71): 'mask_.shape[1] - patch_size'}, {}), '(0, mask_.shape[1] - patch_size)', True, 'import numpy as np\n'), ((50, 23, 50, 72), 'numpy.random.randint', 'np.random.randint', ({(50, 41, 50, 42): '0', (50, 44, 50, 71): 'mask_.shape[0] - patch_size'}, {}), '(0, mask_.shape[0] - patch_size)', True, 'import numpy as np\n'), ((56, 24, 56, 52), 'numpy.moveaxis', 'np.moveaxis', ({(56, 36, 56, 45): 'img_patch', (56, 47, 56, 48): '0', (56, 50, 56, 51): '2'}, {}), '(img_patch, 0, 2)', True, 'import numpy as np\n'), ((69, 25, 69, 54), 'numpy.expand_dims', 'np.expand_dims', ({(69, 40, 69, 50): 'mask_patch', (69, 52, 69, 53): '2'}, {}), '(mask_patch, 2)', True, 'import numpy as np\n'), ((72, 15, 72, 26), 'numpy.array', 'np.array', ({(72, 24, 72, 25): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((72, 28, 72, 39), 'numpy.array', 'np.array', ({(72, 37, 72, 38): 'y'}, {}), '(y)', True, 'import numpy as np\n')]