data.py 文件源码

def load_augment(fname, preprocessor, w, h, is_training, aug_params=no_augmentation_params, transform=None, bbox=None,
                 fill_mode='constant', fill_mode_cval=0, standardizer=None, save_to_dir=None):
    """Load augmented image with output shape (h, w, c)

    Default arguments return non augmented image of shape (h, w, c).
    To apply a fixed transform (and color augmentation) specify transform (and color_vec in standardizer).
    To generate a random augmentation specify aug_params (and sigma in standardizer).
    """
    img = _load_image_th(fname, preprocessor)
    # img shape - (c, h, w)

    if bbox is not None:
        img = _definite_crop(img, bbox)
        # print(img.shape)
        # import cv2
        # cv2.imshow("test", np.asarray(img[1,:,:], dtype=np.uint8))
        # cv2.waitKey(0)
        if bbox[4] == 1:
            img = img[:, :, ::-1]
    elif transform is not None:
        img = _perturb_fixed(img, tform_augment=transform, target_shape=(w, h), mode=fill_mode,
                             mode_cval=fill_mode_cval)
    else:
        img = _perturb(img, augmentation_params=aug_params, target_shape=(w, h), mode=fill_mode,
                       mode_cval=fill_mode_cval)

    if save_to_dir is not None:
        file_full_name = os.path.basename(fname)
        file_name, file_ext = os.path.splitext(file_full_name)
        fname2 = "%s/%s_DA_%d%s" % (save_to_dir, file_name, np.random.randint(1e4), file_ext)
        _save_image_th(img, fname2)

    if standardizer is not None:
        img = standardizer(img, is_training)

    # convert to shape (h, w, c)
    return img.transpose(1, 2, 0)