diff --git a/src/av2/torch/structures/flow.py b/src/av2/torch/structures/flow.py index 13009e13..502d8164 100644 --- a/src/av2/torch/structures/flow.py +++ b/src/av2/torch/structures/flow.py @@ -8,6 +8,8 @@ import torch from kornia.geometry.linalg import transform_points from torch import BoolTensor, ByteTensor, FloatTensor +from copy import deepcopy +import numpy as np from av2.evaluation.scene_flow.constants import ( CATEGORY_TO_INDEX, @@ -82,29 +84,59 @@ def from_sweep_pair(cls, sweeps: Tuple[Sweep, Sweep]) -> Flow: is_valid = torch.ones(len(current_pc), dtype=torch.bool) category_inds = torch.zeros(len(current_pc), dtype=torch.uint8) + # # old version: fixed bounding box expansion + # for id in current_cuboid_map: + # c0 = current_cuboid_map[id] + # c0.length_m += BOUNDING_BOX_EXPANSION # the bounding boxes are a little too tight sometimes + # c0.width_m += BOUNDING_BOX_EXPANSION + # obj_pts_npy, obj_mask_npy = c0.compute_interior_points(current_pc.numpy()) + # obj_pts, obj_mask = torch.as_tensor( + # obj_pts_npy, dtype=torch.float32 + # ), torch.as_tensor(obj_mask_npy) + # category_inds[obj_mask] = CATEGORY_TO_INDEX[str(c0.category)] + + # if id in next_cuboid_map: + # c1 = next_cuboid_map[id] + # c1_SE3_c0 = c1.dst_SE3_object.compose(c0.dst_SE3_object.inverse()) + # flow[obj_mask] = ( + # torch.as_tensor( + # c1_SE3_c0.transform_point_cloud(obj_pts.numpy()), + # dtype=torch.float32, + # ) + # - obj_pts + # ) + # else: + # is_valid[obj_mask] = 0 + + # NOTE(HiMo): box expansion based on the object velocity + # check more detail: https://kin-zhang.github.io/HiMo + current_pc_npy = current_pc.numpy() for id in current_cuboid_map: - c0 = current_cuboid_map[id] - c0.length_m += BOUNDING_BOX_EXPANSION # the bounding boxes are a little too tight sometimes - c0.width_m += BOUNDING_BOX_EXPANSION - obj_pts_npy, obj_mask_npy = c0.compute_interior_points(current_pc.numpy()) - obj_pts, obj_mask = ( - torch.as_tensor(obj_pts_npy, dtype=torch.float32), - torch.as_tensor(obj_mask_npy), - ) - category_inds[obj_mask] = CATEGORY_TO_INDEX[str(c0.category)] - + c0 = deepcopy(current_cuboid_map[id]) + obj_pts_npy, obj_mask_npy = c0.compute_interior_points(current_pc_npy) if id in next_cuboid_map: c1 = next_cuboid_map[id] c1_SE3_c0 = c1.dst_SE3_object.compose(c0.dst_SE3_object.inverse()) + rel_obj_flow = c1_SE3_c0.transform_point_cloud(obj_pts_npy) - obj_pts_npy + delta_move = abs(np.linalg.norm(rel_obj_flow, axis=0).mean()) + if delta_move > 0.04: # only when it's moving 0.4m/s + c0 = current_cuboid_map[id] + c0.length_m += (BOUNDING_BOX_EXPANSION + min(delta_move/2, 2)) # since 180/360 for two LiDARs orientation + c0.width_m += BOUNDING_BOX_EXPANSION + c0.height_m += BOUNDING_BOX_EXPANSION + obj_pts, obj_mask = c0.compute_interior_points(current_pc_npy) + + # after expansion, we need to recompute the flow + c1_SE3_c0 = c1.dst_SE3_object.compose(c0.dst_SE3_object.inverse()) flow[obj_mask] = ( torch.as_tensor( - c1_SE3_c0.transform_point_cloud(obj_pts_npy), + c1_SE3_c0.transform_point_cloud(obj_pts), dtype=torch.float32, ) - obj_pts ) else: - is_valid[obj_mask] = 0 + is_valid[torch.as_tensor(obj_mask_npy)] = 0 dynamic_norm = torch.linalg.vector_norm(flow - rigid_flow, dim=-1) is_dynamic: BoolTensor = dynamic_norm >= SCENE_FLOW_DYNAMIC_THRESHOLD diff --git a/tests/unit/test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede/flow_labels.feather b/tests/unit/test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede/flow_labels.feather index 7df052a0..877cacd1 100644 Binary files a/tests/unit/test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede/flow_labels.feather and b/tests/unit/test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede/flow_labels.feather differ