| 13 | |
| 14 | |
| 15 | class ShapeAutoEncoder(object): |
| 16 | def __init__(self, log_dir, args): |
| 17 | self.log_dir = log_dir |
| 18 | |
| 19 | self.batch_size = args.enc_batch_size |
| 20 | self.n_iters = args.enc_n_iters |
| 21 | self.vol_ratio = args.vol_ratio |
| 22 | |
| 23 | self.fm_reso = args.fm_reso if hasattr(args, "fm_reso") else 128 |
| 24 | |
| 25 | self.data_type = args.data_type |
| 26 | self.sdf_loss_type = args.sdf_loss |
| 27 | self.tex_loss_type = args.tex_loss |
| 28 | |
| 29 | self.device = dist_util.dev() |
| 30 | self.tex_weight = args.tex_weight |
| 31 | self.tex_threshold_ratio = args.tex_threshold_ratio |
| 32 | |
| 33 | self.sdf_renorm = args.sdf_renorm |
| 34 | |
| 35 | self.init_lr = args.enc_lr |
| 36 | self.lr_split = args.enc_lr_split if hasattr(args, "enc_lr_split") else -1 |
| 37 | self.min_lr_ratio = args.enc_lr_decay if hasattr(args, "enc_lr_decay") else 0.01 |
| 38 | |
| 39 | # build network |
| 40 | self.net = get_networks(args).to(self.device) |
| 41 | |
| 42 | # input data information |
| 43 | self.Ka = None |
| 44 | self.Kd = None |
| 45 | self.Ks = None |
| 46 | self.Ns = None |
| 47 | self.aabb = None |
| 48 | self.featmap_size = None |
| 49 | self.sdf_threshold = None |
| 50 | |
| 51 | def _load_data(self, path, sdf_renorm=False): |
| 52 | data = np.load(path) |
| 53 | self.aabb = torch.from_numpy(data["aabb"]).float().to(self.device) |
| 54 | self.sdf_threshold = float(data["threshold"]) |
| 55 | self.Ka = data["Ka"].tolist() if "Ka" in data else [0, 0, 0] |
| 56 | self.Kd = data["Kd"].tolist() if "Kd" in data else [1, 1, 1] |
| 57 | self.Ks = data["Ks"].tolist() if "Ks" in data else [0.4, 0.4, 0.4] |
| 58 | self.Ns = data["Ns"].tolist() if "Ns" in data else 10 |
| 59 | print("aabb: ", self.aabb) |
| 60 | print("using sdf_threshold: ", self.sdf_threshold) |
| 61 | |
| 62 | pts_grid = data["pts_grid"] |
| 63 | sdf_gird = data["sdf_grid"] |
| 64 | pts_near_surf = data["pts_near_surf"] |
| 65 | sdf_near_surf = data["sdf_near_surf"] |
| 66 | |
| 67 | if self.data_type != "sdf": |
| 68 | tex_grid = data["tex_grid"] |
| 69 | pts_on_surf = data["pts_on_surf"] |
| 70 | tex_on_surf = data["tex_on_surf"] |
| 71 | tex_near_surf = data["tex_near_surf"] |
| 72 | print("pts_grid shape: ", pts_grid.shape) |