diff --git a/config/terrain_estimation/voxel_to_bev_8.yaml b/config/terrain_estimation/voxel_to_bev_8.yaml index 95ac36d..1933cb8 100644 --- a/config/terrain_estimation/voxel_to_bev_8.yaml +++ b/config/terrain_estimation/voxel_to_bev_8.yaml @@ -61,6 +61,19 @@ # kernel_radius: 0.5 # kernel_sharpness: 1.0 +- + type: roughness + args: + # input_layer: min_elevation_filtered_inflated + input_layer: terrain + mask_layer: min_elevation_filtered_mask + + thresh: 0.05 #at least this frac of neighboring cells in the kernel must be observed + + kernel_params: + kernel_type: gaussian + + kernel_radius: 2.0 #kernel radius in m - type: slope diff --git a/physics_atv_visual_mapping/terrain_estimation/processing_blocks/roughness.py b/physics_atv_visual_mapping/terrain_estimation/processing_blocks/roughness.py new file mode 100644 index 0000000..affdcf4 --- /dev/null +++ b/physics_atv_visual_mapping/terrain_estimation/processing_blocks/roughness.py @@ -0,0 +1,61 @@ +import torch +import torch_scatter + +from physics_atv_visual_mapping.terrain_estimation.processing_blocks.base import TerrainEstimationBlock +from physics_atv_visual_mapping.terrain_estimation.processing_blocks.utils import setup_kernel, apply_kernel +from physics_atv_visual_mapping.feature_key_list import FeatureKeyList + +class Roughness(TerrainEstimationBlock): + """ + Compute a per-cell min and max height + """ + def __init__(self, voxel_metadata, voxel_feature_keys, input_layer, mask_layer, kernel_params, thresh, device): + """ + Args: + thresh: at least this fraction of neighboring cells must be observed + """ + super().__init__(voxel_metadata, voxel_feature_keys, device) + self.input_layer = input_layer + self.mask_layer = mask_layer + + self.kernel = setup_kernel(**kernel_params, metadata=voxel_metadata).to(device) + self.thresh = thresh * self.kernel.sum() + + def to(self, device): + self.device = device + return self + + @property + def output_feature_keys(self): + return FeatureKeyList( + label = ["roughness"], + metainfo = ["terrain_estimation"] + ) + + def run(self, voxel_grid, bev_grid): + input_idx = bev_grid.feature_keys.index(self.input_layer) + input_data = bev_grid.data[..., input_idx].clone() + + mask_idx = bev_grid.feature_keys.index(self.mask_layer) + valid_mask = bev_grid.data[..., mask_idx] > 1e-4 + + input_data[~valid_mask] = 0. + input_data_sq = torch.pow(input_data, 2) + + sum_x = apply_kernel(kernel=self.kernel, data=input_data) + sum_x2 = apply_kernel(kernel=self.kernel, data=input_data_sq) + count = apply_kernel(kernel=self.kernel, data=valid_mask.float()) + + roughness = torch.zeros_like(input_data) + + # var = E[x^2] - E[x]^2 + E_x2 = sum_x2[valid_mask] / count[valid_mask] + E_x = sum_x[valid_mask] / count[valid_mask] + + roughness[valid_mask] = torch.sqrt(E_x2 - (torch.pow(E_x, 2))) + + output_data_idx = bev_grid.feature_keys.index(self.output_feature_keys.label[0]) + + bev_grid.data[..., output_data_idx] = roughness + + return bev_grid diff --git a/physics_atv_visual_mapping/terrain_estimation/terrain_estimation_pipeline.py b/physics_atv_visual_mapping/terrain_estimation/terrain_estimation_pipeline.py index 513abca..6acac4f 100644 --- a/physics_atv_visual_mapping/terrain_estimation/terrain_estimation_pipeline.py +++ b/physics_atv_visual_mapping/terrain_estimation/terrain_estimation_pipeline.py @@ -11,6 +11,7 @@ from physics_atv_visual_mapping.terrain_estimation.processing_blocks.terrain_inflation import TerrainInflation from physics_atv_visual_mapping.terrain_estimation.processing_blocks.mrf_terrain_estimation import MRFTerrainEstimation from physics_atv_visual_mapping.terrain_estimation.processing_blocks.porosity import Porosity +from physics_atv_visual_mapping.terrain_estimation.processing_blocks.roughness import Roughness from physics_atv_visual_mapping.terrain_estimation.processing_blocks.sdf import SDF from physics_atv_visual_mapping.terrain_estimation.processing_blocks.slope import Slope from physics_atv_visual_mapping.terrain_estimation.processing_blocks.terrain_normals_gradient import TerrainNormalsGradient @@ -44,6 +45,8 @@ def setup_terrain_estimation_pipeline(config, voxel_grid): block = MRFTerrainEstimation(**block_config["args"]) elif btype == "porosity": block = Porosity(**block_config["args"]) + elif btype == "roughness": + block = Roughness(**block_config["args"]) elif btype == "slope": block = Slope(**block_config["args"]) elif btype == "normals_sobel":