def test_cubemap_stiching(test_cfg_path: str, mode: str, camera: str):
meta_config = get_config(config_paths=test_cfg_path)
meta_config.defrost()
config = meta_config.TASK_CONFIG
CAMERA_NUM = 6
orient = [
[0, math.pi, 0], # Back
[-math.pi / 2, 0, 0], # Down
[0, 0, 0], # Front
[0, math.pi / 2, 0], # Right
[0, 3 / 2 * math.pi, 0], # Left
[math.pi / 2, 0, 0], # Up
]
sensor_uuids = []
if "RGB_SENSOR" in config.SIMULATOR.AGENT_0.SENSORS:
config.SIMULATOR.RGB_SENSOR.ORIENTATION = orient[0]
for camera_id in range(1, CAMERA_NUM):
camera_template = f"RGB_{camera_id}"
camera_config = deepcopy(config.SIMULATOR.RGB_SENSOR)
camera_config.ORIENTATION = orient[camera_id]
camera_config.UUID = camera_template.lower()
sensor_uuids.append(camera_config.UUID)
setattr(config.SIMULATOR, camera_template, camera_config)
config.SIMULATOR.AGENT_0.SENSORS.append(camera_template)
if "DEPTH_SENSOR" in config.SIMULATOR.AGENT_0.SENSORS:
config.SIMULATOR.DEPTH_SENSOR.ORIENTATION = orient[0]
for camera_id in range(1, CAMERA_NUM):
camera_template = f"DEPTH_{camera_id}"
camera_config = deepcopy(config.SIMULATOR.DEPTH_SENSOR)
camera_config.ORIENTATION = orient[camera_id]
camera_config.UUID = camera_template.lower()
sensor_uuids.append(camera_config.UUID)
setattr(config.SIMULATOR, camera_template, camera_config)
config.SIMULATOR.AGENT_0.SENSORS.append(camera_template)
meta_config.TASK_CONFIG = config
meta_config.SENSORS = config.SIMULATOR.AGENT_0.SENSORS
if camera == "equirec":
meta_config.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ.SENSOR_UUIDS = tuple(
sensor_uuids)
elif camera == "fisheye":
meta_config.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.SENSOR_UUIDS = tuple(
sensor_uuids)
meta_config.freeze()
execute_exp(meta_config, mode)
# Deinit processes group
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
def test_cubemap_stiching(test_cfg_path: str, mode: str, camera: str,
sensor_type: str):
meta_config = get_config(config_paths=test_cfg_path)
meta_config.defrost()
config = meta_config.TASK_CONFIG
CAMERA_NUM = 6
orient = [
[0, math.pi, 0], # Back
[-math.pi / 2, 0, 0], # Down
[0, 0, 0], # Front
[0, math.pi / 2, 0], # Right
[0, 3 / 2 * math.pi, 0], # Left
[math.pi / 2, 0, 0], # Up
]
sensor_uuids = []
if f"{sensor_type}_SENSOR" not in config.SIMULATOR.AGENT_0.SENSORS:
config.SIMULATOR.AGENT_0.SENSORS.append(f"{sensor_type}_SENSOR")
sensor = getattr(config.SIMULATOR, f"{sensor_type}_SENSOR")
for camera_id in range(CAMERA_NUM):
camera_template = f"{sensor_type}_{camera_id}"
camera_config = deepcopy(sensor)
camera_config.ORIENTATION = orient[camera_id]
camera_config.UUID = camera_template.lower()
sensor_uuids.append(camera_config.UUID)
setattr(config.SIMULATOR, camera_template, camera_config)
config.SIMULATOR.AGENT_0.SENSORS.append(camera_template)
meta_config.TASK_CONFIG = config
meta_config.SENSORS = config.SIMULATOR.AGENT_0.SENSORS
if camera == "equirect":
meta_config.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ.SENSOR_UUIDS = tuple(
sensor_uuids)
elif camera == "fisheye":
meta_config.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.SENSOR_UUIDS = tuple(
sensor_uuids)
meta_config.freeze()
if camera in ["equirect", "fisheye"]:
execute_exp(meta_config, mode)
# Deinit processes group
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
elif camera == "cubemap":
# 1) Generate an equirect image from cubemap images.
# 2) Generate cubemap images from the equirect image.
# 3) Compare the input and output cubemap
env_fn_args = []
for split in ["train", "val"]:
tmp_config = config.clone()
tmp_config.defrost()
tmp_config.DATASET["SPLIT"] = split
tmp_config.freeze()
env_fn_args.append((tmp_config, None))
with VectorEnv(env_fn_args=env_fn_args) as envs:
observations = envs.reset()
batch = batch_obs(observations)
orig_batch = deepcopy(batch)
# ProjectionTransformer
obs_trans_to_eq = baseline_registry.get_obs_transformer(
"CubeMap2Equirect")
cube2equirect = obs_trans_to_eq(sensor_uuids, (256, 512))
obs_trans_to_cube = baseline_registry.get_obs_transformer(
"Equirect2CubeMap")
equirect2cube = obs_trans_to_cube(cube2equirect.target_uuids,
(256, 256))
# Cubemap to Equirect to Cubemap
batch_eq = cube2equirect(batch)
batch_cube = equirect2cube(batch_eq)
# Extract input and output cubemap
output_cube = batch_cube[cube2equirect.target_uuids[0]]
input_cube = [orig_batch[key] for key in sensor_uuids]
input_cube = torch.stack(input_cube, axis=1)
input_cube = torch.flatten(input_cube, end_dim=1)
# Apply blur to absorb difference (blur, etc.) caused by conversion
if sensor_type == "RGB":
output_cube = output_cube.float() / 255
input_cube = input_cube.float() / 255
output_cube = output_cube.permute((0, 3, 1, 2)) # NHWC => NCHW
input_cube = input_cube.permute((0, 3, 1, 2)) # NHWC => NCHW
apply_blur = torch.nn.AvgPool2d(5, 3, 2)
output_cube = apply_blur(output_cube)
input_cube = apply_blur(input_cube)
# Calculate the difference
diff = torch.abs(output_cube - input_cube)
assert diff.mean().item() < 0.01
else:
raise ValueError(f"Unknown camera name: {camera}")