Spaces:
Running
on
Zero
Running
on
Zero
| import numpy as np | |
| from modules.bbox_gen.utils.mesh import change_pcd_range | |
| class BoundsTokenizerDiag: | |
| def __init__(self, bins, BOS_id, EOS_id, PAD_id): | |
| self.bins = bins | |
| self.BOS_id = BOS_id | |
| self.EOS_id = EOS_id | |
| self.PAD_id = PAD_id | |
| def encode(self, data_dict, coord_rg=(-1,1)): | |
| """ | |
| Encode bounding boxes to token sequence | |
| Args: | |
| data_dict: dictionary containing bounding boxes | |
| coord_rg: range of coordinate values | |
| Returns: | |
| token sequence | |
| """ | |
| bounds = data_dict["bounds"] # (s, 2, 3) | |
| all_vertices = bounds.reshape(-1, 6) | |
| all_vertices = change_pcd_range(all_vertices, from_rg=coord_rg, to_rg=(0.5/self.bins, 1-0.5/self.bins)) | |
| quantized_vertices = (all_vertices * self.bins).astype(np.int32) | |
| tokens = [] | |
| tokens.append(self.BOS_id) | |
| tokens.extend(quantized_vertices.flatten().tolist()) | |
| tokens.append(self.EOS_id) | |
| tokens = np.array(tokens) | |
| return tokens | |
| def decode(self, tokens, coord_rg=(-1,1)): | |
| """ | |
| Decode token sequence back to bounding boxes | |
| Args: | |
| tokens: token sequence | |
| Returns: | |
| bounding box array [N, 2, 3] | |
| """ | |
| # Remove special tokens | |
| valid_tokens = [] | |
| for t in tokens: | |
| if t != self.BOS_id and t != self.EOS_id and t != self.PAD_id: | |
| valid_tokens.append(t) | |
| # Ensure correct number of tokens (2 vertices per box, 3 coordinates per vertex) | |
| if len(valid_tokens) % (2 * 3) != 0: | |
| raise ValueError(f"Invalid token count: {len(valid_tokens)}") | |
| # Reshape to vertex coordinates | |
| points = np.array(valid_tokens).reshape(-1, 2, 3) | |
| # Convert quantized coordinates back to continuous values | |
| points = points / self.bins | |
| points = change_pcd_range(points, from_rg=(0.5/self.bins, 1-0.5/self.bins), to_rg=coord_rg) | |
| return points |