(batch)
| 4 | import random |
| 5 | |
| 6 | def my_collate(batch): |
| 7 | src_pc = [] |
| 8 | src_name = [] |
| 9 | key_pts = [] |
| 10 | tar_pc = [] |
| 11 | tar_name = [] |
| 12 | w_pc = [] |
| 13 | src_ver = [] |
| 14 | src_face = [] |
| 15 | tar_ver = [] |
| 16 | tar_face = [] |
| 17 | real_ver = [] |
| 18 | real_face = [] |
| 19 | w_mesh = [] |
| 20 | for data in batch: |
| 21 | src_pc.append(torch.from_numpy(data["src_pc"]).unsqueeze(0).float()) |
| 22 | src_name.append(data["src_name"]) |
| 23 | tar_pc.append(torch.from_numpy(data["tar_pc"]).unsqueeze(0).float()) |
| 24 | tar_name.append(data["tar_name"]) |
| 25 | key_pts.append(torch.from_numpy(data["key_pts"]).unsqueeze(0).float()) |
| 26 | w_pc.append(torch.from_numpy(data["w_pc"]).unsqueeze(0).float()) |
| 27 | src_ver.append(torch.from_numpy(data["src_ver"])) |
| 28 | src_face.append(torch.from_numpy(data["src_face"])) |
| 29 | tar_ver.append(torch.from_numpy(data["tar_ver"])) |
| 30 | tar_face.append(torch.from_numpy(data["tar_face"])) |
| 31 | real_ver.append(torch.from_numpy(data["real_ver"])) |
| 32 | real_face.append(torch.from_numpy(data["real_face"])) |
| 33 | w_mesh.append(torch.from_numpy(data["w_mesh"]).float()) |
| 34 | src_pc = torch.cat(src_pc).cuda() |
| 35 | tar_pc = torch.cat(tar_pc).cuda() |
| 36 | key_pts = torch.cat(key_pts).cuda() |
| 37 | w_pc = torch.cat(w_pc).cuda() |
| 38 | return {"src_pc": src_pc, "src_name": src_name, "key_pts": key_pts, |
| 39 | "tar_pc": tar_pc, "tar_name": tar_name, "w_pc": w_pc, |
| 40 | "src_ver": src_ver, "src_face": src_face, "w_mesh": w_mesh, |
| 41 | "tar_ver": tar_ver, "tar_face": tar_face, |
| 42 | "real_ver": real_ver, "real_face": real_face} |
| 43 | |
| 44 | |
| 45 | class ChairDataset(torch.utils.data.Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected