| 43 | |
| 44 | |
| 45 | class ChairDataset(torch.utils.data.Dataset): |
| 46 | def __init__(self, phase="train", data_dir="../data/chair"): |
| 47 | super().__init__() |
| 48 | self.data_dir = data_dir |
| 49 | self.phase = phase |
| 50 | |
| 51 | with open(os.path.join(self.data_dir, "all.txt")) as f: |
| 52 | lines = f.readlines() |
| 53 | self.models = [line.rstrip() for line in lines] |
| 54 | |
| 55 | self.k = 20 |
| 56 | |
| 57 | with open(os.path.join(self.data_dir, "%s.txt" % self.phase)) as f: |
| 58 | lines = f.readlines() |
| 59 | self.ids = [int(line.rstrip()) for line in lines] |
| 60 | |
| 61 | self.pc = np.load(os.path.join( |
| 62 | self.data_dir, "pc_4096.npy"), allow_pickle=True) |
| 63 | self.key_pts = np.load(os.path.join( |
| 64 | self.data_dir, "key_point_50.npy"), allow_pickle=True) |
| 65 | self.mesh_vertices = np.load(os.path.join( |
| 66 | self.data_dir, "mesh_vertices.npy"), allow_pickle=True) |
| 67 | self.mesh_faces = np.load(os.path.join( |
| 68 | self.data_dir, "mesh_faces.npy"), allow_pickle=True) |
| 69 | # biharmonic weights |
| 70 | self.w_pc = np.load(os.path.join( |
| 71 | self.data_dir, "w_pc_4096.npy"), allow_pickle=True) |
| 72 | w_mesh_0 = np.load(os.path.join( |
| 73 | self.data_dir, "w_mesh_0.npy"), allow_pickle=True) |
| 74 | w_mesh_1 = np.load(os.path.join( |
| 75 | self.data_dir, "w_mesh_1.npy"), allow_pickle=True) |
| 76 | w_mesh_2 = np.load(os.path.join( |
| 77 | self.data_dir, "w_mesh_2.npy"), allow_pickle=True) |
| 78 | w_mesh_3 = np.load(os.path.join( |
| 79 | self.data_dir, "w_mesh_3.npy"), allow_pickle=True) |
| 80 | self.w_mesh = list(w_mesh_0) + list(w_mesh_1) + \ |
| 81 | list(w_mesh_2) + list(w_mesh_3) |
| 82 | |
| 83 | def __len__(self): |
| 84 | return len(self.ids) * self.k |
| 85 | |
| 86 | def __getitem__(self, idx): |
| 87 | src_id = self.ids[idx // self.k] |
| 88 | tar_id = random.choice(self.ids) |
| 89 | # positive sample for discriminator network |
| 90 | real_id = random.choice(self.ids) |
| 91 | src_name = self.models[src_id] |
| 92 | tar_name = self.models[tar_id] |
| 93 | src_pc = self.pc[src_id] |
| 94 | tar_pc = self.pc[tar_id] |
| 95 | key_pts = self.key_pts[src_id] |
| 96 | w_mesh = self.w_mesh[src_id] |
| 97 | w_pc = self.w_pc[src_id] |
| 98 | src_ver = self.mesh_vertices[src_id] |
| 99 | src_face = self.mesh_faces[src_id] |
| 100 | tar_ver = self.mesh_vertices[tar_id] |
| 101 | tar_face = self.mesh_faces[tar_id] |
| 102 | real_ver = self.mesh_vertices[real_id] |
nothing calls this directly
no outgoing calls
no test coverage detected