utils.py 文件源码

def load_image(im_id: str, rgb_only=False, align=True) -> np.ndarray:
    im_rgb = tiff.imread('./three_band/{}.tif'.format(im_id)).transpose([1, 2, 0])
    if rgb_only:
        return im_rgb
    im_p = np.expand_dims(tiff.imread('sixteen_band/{}_P.tif'.format(im_id)), 2)
    im_m = tiff.imread('sixteen_band/{}_M.tif'.format(im_id)).transpose([1, 2, 0])
    im_a = tiff.imread('sixteen_band/{}_A.tif'.format(im_id)).transpose([1, 2, 0])
    w, h = im_rgb.shape[:2]
    if align:
        key = lambda x: '{}_{}'.format(im_id, x)
        im_p, _ = _aligned(im_rgb, im_p, key=key('p'))
        im_m, aligned = _aligned(im_rgb, im_m, im_m[:, :, :3], key=key('m'))
        im_ref = im_m[:, :, -1] if aligned else im_rgb[:, :, 0]
        im_a, _ = _aligned(im_ref, im_a, im_a[:, :, 0], key=key('a'))
    if im_p.shape != im_rgb.shape[:2]:
        im_p = cv2.resize(im_p, (h, w), interpolation=cv2.INTER_CUBIC)
    im_p = np.expand_dims(im_p, 2)
    im_m = cv2.resize(im_m, (h, w), interpolation=cv2.INTER_CUBIC)
    im_a = cv2.resize(im_a, (h, w), interpolation=cv2.INTER_CUBIC)
    return np.concatenate([im_rgb, im_p, im_m, im_a], axis=2)