def prepare_s2_counts(self):
osm_file = get_roads(list(self.bbox_epsg4326), ["motorway", "trunk", "primary"], osm_buffer,
self.dirs["osm"],
str(self.bbox_epsg4326).replace(", ", "_")[1:-1].replace(".", "_")
+ "_osm_roads", "EPSG:" + str(self.detections.crs.to_epsg()))
osm_roads = gpd.read_file(osm_file)
# subset to road type of station
station_road_type = [self.station_meta["road_type"]]
if station_road_type[0] == "primary":
station_road_type.append("trunk")
station_road_type.append(station_road_type[0] + "_link")
osm_subset = []
for road_type in np.unique(station_road_type):
subset = osm_roads[osm_roads["osm_value"] == road_type]
if len(subset) > 0:
osm_subset.append(subset)
self.osm_roads = pd.concat(osm_subset)
osm_union = self.osm_roads.unary_union
# subset s2 counts to road type
detections_within = []
for row in self.detections.iterrows():
row = row[1]
if row.geometry.intersects(osm_union):
detections_within.append(row)
self.detections = gpd.GeoDataFrame(detections_within, crs=self.detections.crs)
def get_osm_mask(bbox, crs, reference_arr, lat_lon_dict, dir_out):
osm_file = get_roads(bbox, ["motorway", "trunk", "primary"],
OSM_BUFFER, dir_out,
str(bbox).replace(", ", "_")[1:-1] + "_osm_roads",
str(crs))
osm_vec = gpd.read_file(osm_file)
ref_xr = xr.DataArray(data=reference_arr,
coords=lat_lon_dict,
dims=["lat", "lon"])
osm_raster = rasterize_osm(osm_vec, ref_xr).astype(np.float32)
return osm_raster
def _get_osm_mask(bbox, crs, reference_arr, lat_lon_dict, dir_out):
osm_file = get_roads(
bbox, ["motorway", "trunk", "primary"], osm_buffer, dir_out,
str(bbox).replace(", ", "_").replace("-", "minus")[1:-1] +
"_osm_roads", str(crs), reference_arr)
try:
osm_vec = gpd.read_file(osm_file)
except errors.DriverError:
return np.zeros_like(reference_arr, dtype=np.float16)
else:
ref_xr = xr.DataArray(data=reference_arr,
coords=lat_lon_dict,
dims=["lat", "lon"])
osm_raster = rasterize_osm(osm_vec, ref_xr).astype(np.float16)
osm_raster[osm_raster != 0] = 1
osm_raster[osm_raster == 0] = np.nan
return osm_raster
def subset(d, n_subs, osm_buffer, dir_osm, dir_out):
tgt_crs = "EPSG:4326"
home = os.path.dirname(os.path.dirname(d))
for this_dir in [dir_out, dir_osm]:
if not os.path.exists(this_dir):
os.mkdir(this_dir)
d1 = os.path.join(d, os.listdir(d)[0], "GRANULE")
if os.path.exists(d1):
d2 = os.path.join(os.path.join(d1,
os.listdir(d1)[0]), "IMG_DATA", "R10m")
files = os.listdir(d2)
band_names = np.array([f.split("_")[2] for f in files])
bands = ["B04", "B03", "B02", "B08"]
file_stack = os.path.join(dir_out, os.path.basename(d) + ".tif")
src = rasterio.open(os.path.join(d2, files[0]))
kwargs = src.meta.copy()
src_crs = src.crs
kwargs.update({
"count": len(bands),
"dtype": np.float32,
"driver": "GTiff"
})
if not os.path.exists(file_stack):
with rasterio.open(file_stack, "w", **kwargs) as tgt:
for i, b in enumerate(bands):
fname = files[np.where(band_names == b)[0][0]]
with rasterio.open(os.path.join(d2, fname)) as src:
tgt.write(
rescale_s2(src.read(1)).astype(np.float32), i + 1)
return
a, b = src.transform * [0, 0], src.transform * [src.height, src.width]
src_corners = np.array([a[0], b[1], b[0], a[1]]).flatten()
transformer = Transformer.from_crs(str(src_crs), tgt_crs)
y0, x0 = transformer.transform(src_corners[0], src_corners[3])
y1, x1 = transformer.transform(src_corners[2], src_corners[1])
bbox_epsg4326 = [y1, x0, y0, x1]
fname = os.path.basename(d) + "_.tif"
fname_pure = fname.split(".")[0][0:-1]
gpkg_out = os.path.join(dir_out, fname_pure + ".gpkg")
gdf = gpd.GeoDataFrame({"index": [1]},
geometry=[
geometry.box(src_corners[0], src_corners[1],
src_corners[2], src_corners[3])
],
crs=str(src_crs))
try:
gdf.to_file(gpkg_out, driver="GPKG")
except fiona.errors.DriverError:
pass
# get OSM data and mask data to roads
osm_file = get_roads(bbox_epsg4326, ["motorway", "trunk", "primary"],
osm_buffer, dir_osm, fname_pure + "osm_roads",
str(src_crs))
osm_vec = gpd.read_file(osm_file)
n = n_subs / 2
n = 1 if n < 1 else n
stack = rasterio.open(file_stack, "r")
n_bands = len(bands)
h, w = stack.height, stack.width
data = np.zeros((n_bands, h, w), dtype=np.float32)
for i in range(n_bands):
data[i] = stack.read(i + 1)
stack.close()
tgt_pixels_y, tgt_pixels_x = int(h / n), int(w / n)
src_lat = get_lat(src_corners, h)
src_lon = get_lon(src_corners, w)
dim = range(int(n))
for y in dim:
for x in dim:
y1, y2 = int(y * tgt_pixels_y), int((y + 1) * tgt_pixels_y)
x1, x2 = int(x * tgt_pixels_x), int((x + 1) * tgt_pixels_x)
file_out = os.path.join(
dir_out,
"_".join([fname_pure, "y" + str(y1), "x" + str(x1)]) +
".tif")
if os.path.exists(file_out):
print("Already exists: %s" % file_out)
else:
tgt_lat = src_lat[y1:y2]
tgt_lon = src_lon[x1:x2]
ref_xr = xr.DataArray(data=np.zeros(
(len(tgt_lat), len(tgt_lon))),
coords={
"lat": tgt_lat,
"lon": tgt_lon
},
dims=["lat", "lon"])
osm_raster = rasterize_osm(osm_vec,
ref_xr).astype(np.float32)
osm_raster[osm_raster != 0] = 1
osm_raster[osm_raster == 0] = np.nan
t = stack.transform
transform = Affine(t[0], t[1], tgt_lon[0], t[3], t[4],
tgt_lat[0])
kwargs.update({
"transform": transform,
"height": tgt_pixels_y,
"width": tgt_pixels_x,
"dtype": np.float32
})
with rasterio.open(file_out, "w", **kwargs) as tgt:
for i in range(n_bands):
data_band = rescale_s2(data[i, y1:y2, x1:x2])
tgt.write(
(data_band * osm_raster).astype(np.float32),
i + 1)
os.remove(file_stack)