def crop_raster(raster_img, vector_data):
"""Crop a raster image according to given vector data and
return the cropped version as numpy array"""
vector_crs = rasterio.crs.CRS(vector_data.crs)
if vector_crs != raster_img.meta["crs"]:
vector_data = vector_data.to_crs(raster_img.meta["crs"].data)
mask = rasterio_mask(raster_img, list(vector_data.geometry), crop=False)[0]
return mask
def get_mask(raster_img, vector_data, nan_value=0):
"""Get a mask from a raster for a given vector data.
Args:
raster_img (rasterio dataset object): the rater image to mask
vector_data (iterable polygon data): the labels to mask according to
nan_value (int): the value to fill nan areas with
Returns:
a binary mask (np.array)"""
# check if both have the same crs
# follow the raster data as its easier, faster
# and doesn't involve saving huge new raster data
vector_crs = rasterio.crs.CRS(vector_data.crs)
if vector_crs != raster_img.meta["crs"]:
vector_data = vector_data.to_crs(raster_img.meta["crs"].data)
mask = rasterio_mask(raster_img, list(vector_data.geometry), crop=False)[0]
binary_mask = mask[0, :, :]
binary_mask[np.isnan(binary_mask)] = nan_value
binary_mask[binary_mask > 0] = 1
return binary_mask
def tile(self,
src,
dest_dir=None,
channel_idxs=None,
nodata=None,
alpha=None,
restrict_to_aoi=False,
dest_fname_base=None,
nodata_threshold=None):
"""An object to tile geospatial image strips into smaller pieces.
Arguments
---------
src : :class:`rasterio.io.DatasetReader` or str
The source dataset to tile.
nodata_threshold : float, optional
Nodata percentages greater than this threshold will not be saved as tiles.
restrict_to_aoi : bool, optional
Requires aoi_boundary. Sets all pixel values outside the aoi_boundary to the nodata value of the src image.
"""
src = _check_rasterio_im_load(src)
restricted_im_path = os.path.join(
self.dest_dir, "aoi_restricted_" + os.path.basename(src.name))
self.src_name = src.name # preserves original src name in case restrict is used
if restrict_to_aoi is True:
if self.aoi_boundary is None:
raise ValueError(
"aoi_boundary must be specified when RasterTiler is called."
)
mask_geometry = self.aoi_boundary.intersection(
box(*src.bounds
)) # prevents enlarging raster to size of aoi_boundary
index_lst = list(np.arange(1, src.meta['count'] + 1))
# no need to use transform t since we don't crop. cropping messes up transform of tiled outputs
arr, t = rasterio_mask(src, [mask_geometry],
all_touched=False,
invert=False,
nodata=src.meta['nodata'],
filled=True,
crop=False,
pad=False,
pad_width=0.5,
indexes=list(index_lst))
with rasterio.open(restricted_im_path, 'w', **src.profile) as dest:
dest.write(arr)
dest.close()
src.close()
src = _check_rasterio_im_load(
restricted_im_path
) #if restrict_to_aoi, we overwrite the src to be the masked raster
tile_gen = self.tile_generator(src, dest_dir, channel_idxs, nodata,
alpha, self.aoi_boundary,
restrict_to_aoi)
if self.verbose:
print('Beginning tiling...')
self.tile_paths = []
if nodata_threshold is not None:
if nodata_threshold > 1:
raise ValueError(
"nodata_threshold should be expressed as a float less than 1."
)
print(
"nodata value threshold supplied, filtering based on this percentage."
)
new_tile_bounds = []
for tile_data, mask, profile, tb in tqdm(tile_gen):
nodata_count = np.logical_or.reduce(
(tile_data == profile['nodata']), axis=0).sum()
nodata_perc = nodata_count / (tile_data.shape[1] *
tile_data.shape[2])
if nodata_perc < nodata_threshold:
dest_path = self.save_tile(tile_data, mask, profile,
dest_fname_base)
self.tile_paths.append(dest_path)
new_tile_bounds.append(tb)
else:
print(
"{} of nodata is over the nodata_threshold, tile not saved."
.format(nodata_perc))
self.tile_bounds = new_tile_bounds # only keep the tile bounds that make it past the nodata threshold
else:
for tile_data, mask, profile, tb in tqdm(tile_gen):
dest_path = self.save_tile(tile_data, mask, profile,
dest_fname_base)
self.tile_paths.append(dest_path)
if self.verbose:
print('Tiling complete. Cleaning up...')
self.src.close()
if os.path.exists(os.path.join(self.dest_dir, 'tmp.tif')):
os.remove(os.path.join(self.dest_dir, 'tmp.tif'))
if os.path.exists(restricted_im_path):
os.remove(restricted_im_path)
if self.verbose:
print("Done. CRS returned for vector tiling.")
return _check_crs(
profile['crs']) # returns the crs to be used for vector tiling
clean_value = []
debris_value = []
background_value = []
for image_path in images:
_filename = image_path.split("/")[-1].split(".")[0]
raster_img = rasterio.open(image_path)
img_np = np.moveaxis(raster_img.read(), 0, 2)
img_np[np.isnan(img_np)] = 0
vector_crs = rasterio.crs.CRS(clean.crs)
if vector_crs != raster_img.meta["crs"]:
clean = clean.to_crs(raster_img.meta["crs"].data)
debris = debris.to_crs(raster_img.meta["crs"].data)
clean_mask = rasterio_mask(raster_img,
list(clean.geometry),
crop=False)[0]
clean_mask = clean_mask[0, :, :]
debris_mask = rasterio_mask(raster_img,
list(debris.geometry),
crop=False)[0]
debris_mask = debris_mask[0, :, :]
clean_index = np.argwhere(clean_mask != 0)
debris_index = np.argwhere(debris_mask != 0)
background_index = np.argwhere((clean_mask + debris_mask) == 0)
np.random.shuffle(clean_index)
np.random.shuffle(debris_index)
np.random.shuffle(background_index)
def crop_and_mask(
crop: Polygon,
mask: MultiPolygon,
raster_path: Path,
) -> Dict:
"""
Crop and mask a given raster path.
Supports single band LiDAR rasters and RGBZ rasters.
"""
# Reduce mask size to crop area, all other masks are superflous
mask = crop.intersection(mask)
if isinstance(mask, GeometryCollection):
mask = MultiPolygon([
feature for feature in mask
if not isinstance(feature, (Point, LineString))
])
result = {"shapely_mask": mask, "shapely_crop": crop}
with rasterio.open(raster_path) as src:
assert str(src.crs["proj"]) == "utm" and int(src.crs["zone"]) == 32
original_metadata = src.meta.copy()
bands = src.count
if bands == 1:
# Raster file only contains one band; interpreting as LiDAR data
with_rgb = False
cropped_lidar_data, affine_transformation = rasterio_mask(
dataset=src,
shapes=[crop],
all_touched=True,
crop=True,
filled=False,
)
result["lidar_array"] = cropped_lidar_data
elif bands == 4:
# Raster file contains four bands; interpreting as ZRGB data
with_rgb = True
raster_bands = {1, 2, 3, 4}
lidar_band = 1 + lidar_band_index(raster_path=raster_path, )
rgb_bands = sorted(list(raster_bands - {lidar_band}))
cropped_lidar_data, affine_transformation = rasterio_mask(
dataset=src,
shapes=[crop],
all_touched=True,
crop=True,
filled=False,
indexes=[lidar_band],
)
cropped_aerial_data, affine_aerial_transformation = rasterio_mask(
dataset=src,
shapes=[crop],
all_touched=True,
crop=True,
filled=False,
indexes=rgb_bands,
)
# Aerial data should have same shape as LiDAR data
assert cropped_aerial_data.shape == (
3,
*cropped_lidar_data.shape[1:],
)
# RGB data should be in domain {0, 1, ..., 255}
assert cropped_aerial_data.dtype == "uint8"
# Transformations should be identical
assert (affine_transformation == affine_aerial_transformation)
result["lidar_array"] = cropped_lidar_data
result["rgb_array"] = cropped_aerial_data
else:
raise NotImplementedError(f"Unsupported number of bands = {bands}")
lidar_metadata = original_metadata.copy()
lidar_metadata.update({
"count": 1,
"height": cropped_lidar_data.shape[1],
"width": cropped_lidar_data.shape[2],
"transform": affine_transformation,
"driver": "GTiff",
"dtype": cropped_lidar_data.dtype,
})
cropped_lidar_file = MemoryFile()
with rasterio.open(cropped_lidar_file, "w", **lidar_metadata) as file:
file.write(cropped_lidar_data)
if mask:
mask_data, _ = rasterio_mask(
file,
shapes=[mask],
all_touched=True,
crop=False,
)
mask_data[mask_data > file.nodata] = 1
mask_data[mask_data != 1] = 0
mask_data = mask_data.astype("uint8", copy=False)
else:
mask_data = np.zeros(
cropped_lidar_data.shape,
dtype="uint8",
)
mask_metadata = lidar_metadata.copy()
mask_metadata.update({"dtype": "uint8", "nodata": int(2**8 - 1)})
mask_file = MemoryFile()
with rasterio.open(mask_file, "w", **mask_metadata) as file:
file.write(mask_data)
if with_rgb:
rgb_metadata = lidar_metadata.copy()
rgb_metadata.update({
"dtype": cropped_aerial_data.dtype,
"count": 3,
})
rgb_file = MemoryFile()
with rasterio.open(rgb_file, "w", **rgb_metadata) as file:
file.write(cropped_aerial_data)
result["rgb_file"] = rgb_file
result["lidar_file"] = cropped_lidar_file
result["mask_file"] = mask_file
result["mask_array"] = mask_data
return result