| 45 | |
| 46 | |
| 47 | class MultiTempV2XSimSeg(V2XSimSeg): |
| 48 | def __init__( |
| 49 | self, |
| 50 | dataset_roots=None, |
| 51 | config=None, |
| 52 | split=None, |
| 53 | cache_size=1000, |
| 54 | val=False, |
| 55 | com=False, |
| 56 | bound=None, |
| 57 | kd_flag=False, |
| 58 | no_cross_road=False, |
| 59 | time_stamp = 2, |
| 60 | ): |
| 61 | """ |
| 62 | Inherited from the V2XSimSeg dataset, now support multi timestamp loading |
| 63 | output padded_voxel_points_next are for other timestamp |
| 64 | it is an empty list if timestamp = 1 |
| 65 | """ |
| 66 | super().__init__(dataset_roots, config, split, cache_size, val, com, bound, kd_flag, no_cross_road) |
| 67 | self.time_stamp = config.time_stamp if hasattr(config, 'time_stamp') else time_stamp |
| 68 | print("use time stamp", self.time_stamp) |
| 69 | |
| 70 | def get_seginfo_from_single_agent(self, agent_id, idx): |
| 71 | empty_flag = False |
| 72 | if idx in self.cache[agent_id]: |
| 73 | gt_dict = self.cache[agent_id][idx] |
| 74 | else: |
| 75 | seq_file = self.seq_files[agent_id][idx] |
| 76 | gt_data_handle = np.load(seq_file, allow_pickle=True) |
| 77 | if gt_data_handle == 0: |
| 78 | empty_flag = True |
| 79 | if self.time_stamp > 1: |
| 80 | padded_voxel_next = torch.zeros((self.time_stamp-1, 256, 256, 13)).bool() |
| 81 | else: |
| 82 | padded_voxel_next = torch.zeros(0).bool() |
| 83 | if self.com: |
| 84 | return ( |
| 85 | torch.zeros((256, 256, 13)).bool(), |
| 86 | padded_voxel_next, |
| 87 | torch.zeros((256, 256, 13)).bool(), |
| 88 | torch.zeros((256, 256)).int(), |
| 89 | torch.zeros((self.num_agent, 4, 4)), |
| 90 | 0, |
| 91 | 0, |
| 92 | ) |
| 93 | else: |
| 94 | return ( |
| 95 | torch.zeros((256, 256, 13)).bool(), |
| 96 | padded_voxel_next, |
| 97 | torch.zeros((256, 256, 13)).bool(), |
| 98 | torch.zeros((256, 256)).int(), |
| 99 | ) |
| 100 | else: |
| 101 | gt_dict = gt_data_handle.item() |
| 102 | if len(self.cache[agent_id]) < self.cache_size: |
| 103 | self.cache[agent_id][idx] = gt_dict |
| 104 | |