| 33 | return cat_shape_dict |
| 34 | |
| 35 | class OurDataset(data.Dataset): |
| 36 | |
| 37 | def __init__(self, data_root_dir, data_csv_file, data_features=[], num_points = 1024 ,num_query_points = 1024 ,data_per_seg = 1): |
| 38 | self.data_root_dir = data_root_dir |
| 39 | self.data_csv_file = data_csv_file |
| 40 | self.num_points = num_points |
| 41 | self.num_query_points = num_query_points |
| 42 | self.data_features = data_features |
| 43 | self.data_per_seg = data_per_seg |
| 44 | self.dataset = [] |
| 45 | |
| 46 | with open(self.data_csv_file, 'r') as fin: |
| 47 | self.category_list = [line.strip() for line in fin.readlines()] |
| 48 | |
| 49 | |
| 50 | def transform_pc_to_rot(self, pcs): |
| 51 | # zero-centered |
| 52 | pc_center = (pcs.max(axis=0, keepdims=True) + pcs.min(axis=0, keepdims=True)) / 2 |
| 53 | pc_center = pc_center[0] |
| 54 | new_pcs = pcs - pc_center |
| 55 | |
| 56 | # (batch_size, 2, 3) |
| 57 | def bgs(d6s): |
| 58 | bsz = d6s.shape[0] |
| 59 | b1 = F.normalize(d6s[:, :, 0], p=2, dim=1) |
| 60 | a2 = d6s[:, :, 1] |
| 61 | b2 = F.normalize(a2 - torch.bmm(b1.view(bsz, 1, -1), a2.view(bsz, -1, 1)).view(bsz, 1) * b1, p=2, dim=1) |
| 62 | b3 = torch.cross(b1, b2, dim=1) |
| 63 | return torch.stack([b1, b2, b3], dim=1).permute(0, 2, 1) |
| 64 | |
| 65 | # randomly sample two rotation matrices |
| 66 | rotmat = bgs(torch.rand(1, 6).reshape(-1, 2, 3).permute(0, 2, 1)) |
| 67 | new_pcs = (rotmat.reshape(3, 3) @ new_pcs.T).T |
| 68 | |
| 69 | gt_rot = rotmat[:, :, :2].permute(0, 2, 1).reshape(6).numpy() |
| 70 | |
| 71 | return new_pcs, pc_center, gt_rot |
| 72 | |
| 73 | |
| 74 | def load_data(self): |
| 75 | bar = ProgressBar() |
| 76 | |
| 77 | for category_i in bar(range(len(self.category_list))): |
| 78 | category_id = self.category_list[category_i] |
| 79 | instance_dir = os.path.join(self.data_root_dir, category_id) |
| 80 | fileA = os.path.join(instance_dir, 'partA-pc.csv') |
| 81 | fileB = os.path.join(instance_dir, 'partB-pc.csv') |
| 82 | |
| 83 | if not os.path.exists(fileA) or not os.path.exists(fileB): |
| 84 | print("fileA is", fileA) |
| 85 | print("fileB is", fileB) |
| 86 | print("file not exists") |
| 87 | continue |
| 88 | |
| 89 | dataframe_A = pd.read_csv(fileA, header=None) |
| 90 | dataframe_B = pd.read_csv(fileB, header=None) |
| 91 | gt_pcs_A = dataframe_A.to_numpy() |
| 92 | gt_pcs_B = dataframe_B.to_numpy() |