import cv2
import rasterio
import numpy as np
import tifffile as tiff
from PIL import Image
from rasterio.windows import Window
LOAD_ORIGINAL = True
RGB_ONLY = True
BAND_SWAP = False
[docs]
def pad(img, pad_size=32):
"""
Pad image on the sides, so that eash side is divisible by 32 (network requirement)
if pad = True:
returns image as numpy.array,
tuple with padding in pixels as(x_min_pad, y_min_pad, x_max_pad, y_max_pad)
else:
returns image as numpy.array
"""
if pad_size == 0:
return img
height, width = img.shape[:2]
if height % pad_size == 0:
y_min_pad = 0
y_max_pad = 0
else:
y_pad = pad_size - height % pad_size
y_min_pad = int(y_pad / 2)
y_max_pad = y_pad - y_min_pad
if width % pad_size == 0:
x_min_pad = 0
x_max_pad = 0
else:
x_pad = pad_size - width % pad_size
x_min_pad = int(x_pad / 2)
x_max_pad = x_pad - x_min_pad
img = cv2.copyMakeBorder(img, y_min_pad, y_max_pad,
x_min_pad, x_max_pad, cv2.BORDER_REFLECT_101)
return img, (x_min_pad, y_min_pad, x_max_pad, y_max_pad)
[docs]
def unpad(img, pads):
"""
img: numpy array of the shape (height, width)
pads: (x_min_pad, y_min_pad, x_max_pad, y_max_pad)
@return padded image
"""
(x_min_pad, y_min_pad, x_max_pad, y_max_pad) = pads
height, width = img.shape[:2]
return img[y_min_pad:height - y_max_pad, x_min_pad:width - x_max_pad]
[docs]
def minmax(img):
img_dim = len(img.shape)
if (img_dim < 3):
img = np.expand_dims(img, axis=2)
out = np.zeros_like(img).astype(np.float32)
if img.sum() == 0:
if img_dim < 3:
return out[:, :, 0].astype(np.float32)
else:
return out.astype(np.float32)
# c = img.min()
# d = img.max()
for i in range(img.shape[2]):
c = img[:, :, i].min()
d = img[:, :, i].max()
t = (img[:, :, i] - c) / (float(d - c) + 0.001)
out[:, :, i] = t
if img_dim < 3:
return out[:, :, 0].astype(np.float32)
else:
return out.astype(np.float32)
[docs]
def equalizeRGB(img):
img = minmax(img) * 255
img_y_cr_cb = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2YCrCb)
y, cr, cb = cv2.split(img_y_cr_cb)
# Applying equalize Hist operation on Y channel.
y_eq = cv2.equalizeHist(y)
img_y_cr_cb_eq = cv2.merge((y_eq, cr, cb))
img_rgb_eq = cv2.cvtColor(img_y_cr_cb_eq, cv2.COLOR_YCR_CB2BGR)
return img_rgb_eq
# This is from Antonio
[docs]
def normalizeRGB(img, percent_range=(2, 98)):
min_val, max_val = percent_range
img = img.astype(np.float32)
img_dim = len(img.shape)
if (img_dim < 3):
img = np.expand_dims(img, axis=2)
for b in range(img.shape[2]):
img_b = img[:, :, b]
valid_values = img_b[~np.isnan(img_b)]
if valid_values.size == 0:
min = max = 0
else:
min, max = np.percentile(valid_values, (min_val, max_val))
img[:, :, b] -= min
if min != max:
img[:, :, b] /= (max - min)
img.clip(min=0, max=1, out=img)
if img_dim < 3:
return img[:, :, 0]
else:
return img
[docs]
def load_image(file_name_rgb, file_name_tif):
# pdb.set_trace()
if RGB_ONLY:
if file_name_rgb.lower().endswith(('.png', '.jpg', '.jpeg', 'jp2')):
rgb = np.asarray(Image.open(str(file_name_rgb)))
elif file_name_rgb.lower().endswith(('.tif', '.tiff')):
rgb = tiff.imread(str(file_name_rgb))
# Some of the spacenet data has channel as the 1st dim
#rgb = np.swapaxes(rgb, 0, 2)
#rgb = np.swapaxes(rgb, 0, 1)
#rgb = rgb[:,:,[0,2,1]]
elif file_name_rgb.lower().endswith(('.ntf')):
rgb = readRasterImage(file_name_rgb)[0]
else:
print("Not valid RGB image format")
if rgb.max() > 255:
rgb = (255 * minmax(rgb)).astype(np.uint8)
if len(rgb.shape) < 3:
rgb = np.dstack((rgb, rgb, rgb))
if LOAD_ORIGINAL:
return rgb
if rgb.shape[2] == 4 or RGB_ONLY or not BAND_SWAP:
rgb = rgb[:, :, :3]
else:
rgb = rgb[:, :, [4, 2, 1]]
# if EQUALIZE_FLAG:
# rgb = equalizeRGB(rgb)
# normalize it to [0, 1]
rgb = minmax(rgb)
tf = np.concatenate([rgb, rgb, rgb[:, :, 0:2]], axis=2)
else:
rgb = tiff.imread(str(file_name_tif))
rgb = minmax(rgb)
if len(rgb.shape) < 3:
rgb = np.dstack((rgb, rgb, rgb))
#tf = tiff.imread(str(file_name_tif)).astype(np.float32) / (2 ** 11 - 1)
tf = tiff.imread(str(file_name_tif)).astype(np.float32)
tf = minmax(tf)
if len(tf.shape) < 3:
tf = np.dstack((tf, tf, tf))
if tf.shape[2] == 4:
tf = np.concatenate([tf, tf], axis=2)
if tf.shape[2] == 4 or rgb.shape[2] == 4 or RGB_ONLY or not BAND_SWAP:
return np.concatenate([rgb[:, :, :3], tf], axis=2) * (2 ** 8 - 1)
else:
if tf.shape[2] <= 4:
return np.concatenate([rgb[:, :, :3], tf], axis=2) * (2 ** 8 - 1)
else:
return np.concatenate(
[rgb[:, :, [4, 2, 1]], tf], axis=2) * (2 ** 8 - 1)
[docs]
def readRasterImage(filename: str, geoProps={}, convert=False):
"""Reads ntf files
[description]
Decorators:
task
Arguments:
filename {str} -- filename of the file to read in
Keyword Arguments:
geoProps {dict} -- additional properties, like bounding box (default: {{}})
convert {bool} -- should raster be converted to float and normalized (default: {False})
Returns:
Tuple of (img, transform, bands, datatype)
"""
print('processing %s' % filename)
with rasterio.open(filename, "r") as src_ds:
bands = src_ds.count # num bands
datatype = np.dtype(src_ds.dtypes[0])
transform = src_ds.transform
# --- Initilizing img ---
# Running under the assumption all bands are the same data type
# if bool(geoProps):
# vrt = WarpedVRT(src_ds)
# dst_window = vrt.window( *geoProps["bounds"])
# src_ds = vrt.read(1, window=dst_window)
# img = np.ndarray(
# (*src_ds.shape, bands),
# dtype=datatype)
# else:
img = np.ndarray(
(src_ds.height, src_ds.width, bands),
dtype=datatype)
# --- End: Initilizing img ---
# --- Copying imagery by ROIs ---
# Imagery needs to be cut (cropped?) into ROIs that satisfy
# a 32-pixel size based geometry, due to the network
print(img.shape)
# Creating ROIs
width = 128
shape = img.shape
rois = []
for y in range(0, shape[0], width):
for x in range(0, shape[1], width):
locx = x
locy = y
if x + width >= shape[1]:
locx = shape[1] - width
if y + width >= shape[0]:
locy = shape[0] - width
rois.append((locx, locx + width, locy, locy + width))
# End: Creating ROIs
# Copying by ROI, transposed s.t. the number of bands is the last
# dimension (AD: is this openCV shennanigans?)
for i in range(len(rois)):
win = Window(int(rois[i][0]), int(rois[i][2]), width, width)
for j in range(bands):
# if bool(geoProps):
# data = src_ds[ int(rois[i][0]):int(rois[i][1]), int(rois[i][2]):int(rois[i][3]) ]
# else:
data = src_ds.read(j + 1, window=win)
np.copyto(img[rois[i][2]:rois[i][3],
rois[i][0]:rois[i][1], j], data)
# End: Copying by ROI
if convert:
img = img.astype(np.float32)
if datatype == np.uint8:
img /= 255.0
if bands == 3:
normalizeRGB(img)
else:
# FIXME: THIS IS A NAME ERROR
normalize(img, bands, datatype) # NOQA
# --- End: Copying imagery by ROIs ---
return (img, transform, bands, datatype)