95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
from climatenet import *
|
|
import rasterio
|
|
|
|
|
|
|
|
print("main.py started")
|
|
|
|
|
|
def load_raster(file_path: str):
|
|
with rasterio.open(file_path) as src:
|
|
data = src.read()
|
|
|
|
tensor = torch.from_numpy(data.astype(np.float32))
|
|
return tensor
|
|
|
|
|
|
def save_raster_like(tensor: torch.Tensor, ref_path: str, file_path: str):
|
|
data = tensor.detach().cpu().numpy()
|
|
|
|
with rasterio.open(ref_path) as example:
|
|
meta = example.meta.copy()
|
|
|
|
meta.update({
|
|
"driver": "GTiff",
|
|
"height": data.shape[1],
|
|
"width": data.shape[2],
|
|
"count": data.shape[0],
|
|
"dtype": data.dtype
|
|
})
|
|
|
|
with rasterio.open(file_path, "w", **meta) as dst:
|
|
dst.write(data)
|
|
|
|
|
|
|
|
|
|
building_height = load_raster("data/INPUT/building_height.tif")
|
|
tree_height = load_raster("data/INPUT/tree_height.tif")
|
|
surface_height = load_raster("data/INPUT/zt.tif")
|
|
pavement_type = load_raster("data/INPUT/pavement_type.tif")
|
|
vegetation_type = load_raster("data/INPUT/vegetation_type.tif")
|
|
water_type = load_raster("data/INPUT/water_type.tif")
|
|
|
|
output_pet = load_raster("data/OUTPUT/bio_pet_xy_av_14h.tif")
|
|
output_temp = load_raster("data/OUTPUT/ta_av_h001_1.0m_14h.tif")
|
|
|
|
|
|
|
|
building_height = F.pad(building_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
tree_height = F.pad(tree_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
surface_height = F.pad(surface_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
pavement_type = F.pad(pavement_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
vegetation_type = F.pad(vegetation_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
water_type = F.pad(water_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
|
|
|
|
building_height_coarse = F.interpolate(building_height, size=(64, 64), mode='bilinear', align_corners=False)
|
|
tree_height_coarse = F.interpolate(tree_height, size=(64, 64), mode='bilinear', align_corners=False)
|
|
surface_height_coarse = F.interpolate(surface_height, size=(64, 64), mode='bilinear', align_corners=False)
|
|
pavement_type_coarse = F.interpolate(pavement_type, size=(64, 64), mode='nearest')
|
|
vegetation_type_coarse = F.interpolate(vegetation_type, size=(64, 64), mode='nearest')
|
|
water_type_coarse = F.interpolate(water_type, size=(64, 64), mode='nearest')
|
|
|
|
|
|
print(building_height.shape)
|
|
print(tree_height.shape)
|
|
print(surface_height.shape)
|
|
print(pavement_type.shape)
|
|
print(vegetation_type.shape)
|
|
print(water_type.shape)
|
|
print(output_pet.shape)
|
|
print(output_temp.shape)
|
|
|
|
|
|
input_data_fine = torch.cat([building_height, tree_height, surface_height, pavement_type, vegetation_type, water_type], dim=1)
|
|
print(f"total dim fine {input_data_fine.shape}")
|
|
|
|
|
|
input_data_coarse = torch.cat([building_height_coarse, tree_height_coarse, surface_height_coarse, pavement_type_coarse, vegetation_type_coarse, water_type_coarse], dim=1)
|
|
print(f"total dim coarse {input_data_coarse.shape}")
|
|
|
|
|
|
|
|
model = UrbanClimateUNet();
|
|
model.eval()
|
|
|
|
out = model.forward(input_data_fine, input_data_coarse)
|
|
|
|
output_temp_pred = out.t_air_mean.squeeze(0)
|
|
output_pet_pred = out.pet_mean.squeeze(0)
|
|
|
|
print(f"PET pred shape: {output_pet_pred.shape}");
|
|
print(f"PET pred shape: {output_temp_pred.shape}");
|
|
|
|
save_raster_like(output_temp_pred, "data/OUTPUT/ta_av_h001_1.0m_14h.tif", "data/PRED/temp.tif")
|
|
save_raster_like(output_pet_pred, "data/OUTPUT/bio_pet_xy_av_14h.tif", "data/PRED/pet.tif") |