From bea46114aeda675d0b7a107043791e41199fb404 Mon Sep 17 00:00:00 2001 From: shawticus Date: Mon, 10 Oct 2022 05:32:44 -0700 Subject: [PATCH 1/5] Generalize data paths for any data source --- README.md | 23 ++-- train_3d.py | 46 ++++---- training/dataset.py | 152 +++++++++---------------- training/discriminator_architecture.py | 32 ++---- training/loss.py | 20 +--- training/networks_get3d.py | 4 +- training/sample_camera_distribution.py | 35 ++---- 7 files changed, 105 insertions(+), 207 deletions(-) diff --git a/README.md b/README.md index 70fab28..1cb1b52 100644 --- a/README.md +++ b/README.md @@ -93,21 +93,20 @@ export PYTHONPATH=$PWD:$PYTHONPATH export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ``` -- Train on the unified generator on cars, motorbikes or chair (Improved generator in - Appendix): +- Train the unified generator on cars, motorbikes or chair (Improved generator in Appendix): ```bash -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=40 --data_camera_mode shapenet_car --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=80 --data_camera_mode shapenet_motorbike --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=400 --data_camera_mode shapenet_chair --dmtet_scale 0.8 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/02958343 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=40 --manifest_dir shapenet_car --dmtet_scale 1.0 --one_3d_generator 1 --fp32 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/03790512 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=80 --manifest_dir shapenet_motorbike --dmtet_scale 1.0 --one_3d_generator 1 --fp32 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/03001627 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=400 --manifest_dir shapenet_chair --dmtet_scale 0.8 --one_3d_generator 1 --fp32 0 ``` -- If want to train on seperate generators (main Figure in the paper): +- If want to train on separate generators (main Figure in the paper): ```bash -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=40 --data_camera_mode shapenet_car --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 0 -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=80 --data_camera_mode shapenet_motorbike --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 0 -python train_3d.py --outdir=PATH_TO_LOG --data=PATH_TO_RENDER_IMG --camera_path PATH_TO_RENDER_CAMERA --gpus=8 --batch=32 --gamma=3200 --data_camera_mode shapenet_chair --dmtet_scale 0.8 --use_shapenet_split 1 --one_3d_generator 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/02958343 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=40 --manifest_dir shapenet_car --dmtet_scale 1.0 --one_3d_generator 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/03790512 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=80 --manifest_dir shapenet_motorbike --dmtet_scale 1.0 --one_3d_generator 0 +python train_3d.py --outdir=./logs --data=./shapenet/img/03001627 --camera_path=./shapenet/camera --gpus=8 --batch=32 --gamma=400 --manifest_dir shapenet_chair --dmtet_scale 0.8 --one_3d_generator 0 ``` If want to debug the model first, reduce the number of gpus to 1 and batch size to 4 via: @@ -123,9 +122,9 @@ If want to debug the model first, reduce the number of gpus to 1 and batch size - Inference could operate on a single GPU with 16 GB memory. ```bash -python train_3d.py --outdir=save_inference_results/shapenet_car --gpus=1 --batch=4 --gamma=40 --data_camera_mode shapenet_car --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH -python train_3d.py --outdir=save_inference_results/shapenet_chair --gpus=1 --batch=4 --gamma=40 --data_camera_mode shapenet_chair --dmtet_scale 0.8 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH -python train_3d.py --outdir=save_inference_results/shapenet_motorbike --gpus=1 --batch=4 --gamma=40 --data_camera_mode shapenet_motorbike --dmtet_scale 1.0 --use_shapenet_split 1 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH +python train_3d.py --outdir=save_inference_results/shapenet_car --gpus=1 --batch=4 --gamma=40 --manifest_dir shapenet_car --dmtet_scale 1.0 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH +python train_3d.py --outdir=save_inference_results/shapenet_chair --gpus=1 --batch=4 --gamma=40 --manifest_dir shapenet_chair --dmtet_scale 0.8 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH +python train_3d.py --outdir=save_inference_results/shapenet_motorbike --gpus=1 --batch=4 --gamma=40 --manifest_dir shapenet_motorbike --dmtet_scale 1.0 --one_3d_generator 1 --fp32 0 --inference_vis 1 --resume_pretrain MODEL_PATH ``` - To generate mesh with textures, add one option to the inference diff --git a/train_3d.py b/train_3d.py index f2b87f8..55dede0 100644 --- a/train_3d.py +++ b/train_3d.py @@ -22,10 +22,13 @@ # ---------------------------------------------------------------------------- def subprocess_fn(rank, c, temp_dir): + print('Subprocess...1') dnnlib.util.Logger(file_name=os.path.join(c.run_dir, 'log.txt'), file_mode='a', should_flush=True) + print('Subprocess...2') # Init torch.distributed. if c.num_gpus > 1: + print('Subprocess...3') init_file = os.path.abspath(os.path.join(temp_dir, '.torch_distributed_init')) if os.name == 'nt': init_method = 'file:///' + init_file.replace('\\', '/') @@ -35,9 +38,10 @@ def subprocess_fn(rank, c, temp_dir): init_method = f'file://{init_file}' torch.distributed.init_process_group( backend='nccl', init_method=init_method, rank=rank, world_size=c.num_gpus) - + print('Subprocess...4') # Init torch_utils. sync_device = torch.device('cuda', rank) if c.num_gpus > 1 else None + print('Subprocess...5') training_stats.init_multiprocessing(rank=rank, sync_device=sync_device) if rank != 0: custom_ops.verbosity = 'none' @@ -98,34 +102,27 @@ def launch_training(c, desc, outdir, dry_run): # Launch processes. print('Launching processes...') torch.multiprocessing.set_start_method('spawn', force=True) + print('Launching processes...2') + with tempfile.TemporaryDirectory() as temp_dir: if c.num_gpus == 1: subprocess_fn(rank=0, c=c, temp_dir=temp_dir) else: + print('Launching processes...3') torch.multiprocessing.spawn(fn=subprocess_fn, args=(c, temp_dir), nprocs=c.num_gpus) # ---------------------------------------------------------------------------- def init_dataset_kwargs(data, opt=None): try: - if opt.use_shapenet_split: - dataset_kwargs = dnnlib.EasyDict( - class_name='training.dataset.ImageFolderDataset', - path=data, use_labels=True, max_size=None, xflip=False, - resolution=opt.img_res, - data_camera_mode=opt.data_camera_mode, - add_camera_cond=opt.add_camera_cond, - camera_path=opt.camera_path, - split='test' if opt.inference_vis else 'train', - ) - else: - dataset_kwargs = dnnlib.EasyDict( - class_name='training.dataset.ImageFolderDataset', - path=data, use_labels=True, max_size=None, xflip=False, resolution=opt.img_res, - data_camera_mode=opt.data_camera_mode, - add_camera_cond=opt.add_camera_cond, - camera_path=opt.camera_path, - ) + dataset_kwargs = dnnlib.EasyDict( + class_name='training.dataset.ImageFolderDataset', + path=data, use_labels=True, max_size=None, xflip=False, resolution=opt.img_res, + manifest_dir=opt.manifest_dir, + add_camera_cond=opt.add_camera_cond, + camera_path=opt.camera_path, + split='test' if opt.inference_vis else 'train', + ) dataset_obj = dnnlib.util.construct_class_by_name(**dataset_kwargs) # Subclass of training.dataset.Dataset. dataset_kwargs.camera_path = opt.camera_path dataset_kwargs.resolution = dataset_obj.resolution # Be explicit about resolution. @@ -168,8 +165,7 @@ def parse_comma_separated_list(s): @click.option('--data', help='Path to the Training data Images', metavar='[DIR]', type=str, default='./tmp') @click.option('--camera_path', help='Path to the camera root', metavar='[DIR]', type=str, default='./tmp') @click.option('--img_res', help='The resolution of image', metavar='INT', type=click.IntRange(min=1), default=1024) -@click.option('--data_camera_mode', help='The type of dataset we are using', type=str, default='shapenet_car', show_default=True) -@click.option('--use_shapenet_split', help='whether use the training split or all the data for training', metavar='BOOL', type=bool, default=False, show_default=False) +@click.option('--manifest_dir', help='Directory containing train.txt, test.txt and val.txt', type=str, default='shapenet_car', show_default=True) ### Configs for 3D generator########## @click.option('--use_style_mixing', help='whether use style mixing for generation during inference', metavar='BOOL', type=bool, default=True, show_default=False) @click.option('--one_3d_generator', help='whether we detach the gradient for empty object', metavar='BOOL', type=bool, default=True, show_default=True) @@ -240,8 +236,8 @@ def main(**kwargs): c.training_set_kwargs, dataset_name = init_dataset_kwargs(data=opts.data, opt=opts) if opts.cond and not c.training_set_kwargs.use_labels: raise click.ClickException('--cond=True requires labels specified in dataset.json') - c.training_set_kwargs.split = 'train' if opts.use_shapenet_split else 'all' - if opts.use_shapenet_split and opts.inference_vis: + c.training_set_kwargs.split = 'train' + if opts.inference_vis: c.training_set_kwargs.split = 'test' c.training_set_kwargs.use_labels = opts.cond c.training_set_kwargs.xflip = False @@ -260,7 +256,7 @@ def main(**kwargs): c.G_kwargs.render_type = opts.render_type c.G_kwargs.use_tri_plane = opts.use_tri_plane - c.D_kwargs.data_camera_mode = opts.data_camera_mode + c.D_kwargs.manifest_dir = opts.manifest_dir c.D_kwargs.add_camera_cond = opts.add_camera_cond c.G_kwargs.tet_res = opts.tet_res @@ -270,7 +266,7 @@ def main(**kwargs): c.batch_size = opts.batch c.batch_gpu = opts.batch_gpu or opts.batch // opts.gpus # c.G_kwargs.geo_pos_enc = opts.geo_pos_enc - c.G_kwargs.data_camera_mode = opts.data_camera_mode + c.G_kwargs.manifest_dir = opts.manifest_dir c.G_kwargs.channel_base = c.D_kwargs.channel_base = opts.cbase c.G_kwargs.channel_max = c.D_kwargs.channel_max = opts.cmax diff --git a/training/dataset.py b/training/dataset.py index 8cf6b0e..d47b14a 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -154,12 +154,12 @@ def __init__( path, # Path to directory or zip. camera_path, # Path to camera resolution=None, # Ensure specific resolution, None = highest available. - data_camera_mode='shapenet_car', + manifest_dir='shapenet_car', add_camera_cond=False, - split='all', + split='train', **super_kwargs # Additional arguments for the Dataset base class. ): - self.data_camera_mode = data_camera_mode + self.manifest_dir = manifest_dir self._path = path self._zipfile = None self.root = path @@ -167,67 +167,32 @@ def __init__( self.add_camera_cond = add_camera_cond root = self._path self.camera_root = camera_path - if data_camera_mode == 'shapenet_car' or data_camera_mode == 'shapenet_chair' \ - or data_camera_mode == 'renderpeople' or data_camera_mode == 'shapenet_motorbike' \ - or data_camera_mode == 'ts_house' \ - or data_camera_mode == 'ts_animal': - print('==> use shapenet dataset') - if not os.path.exists(root): - print('==> ERROR!!!! THIS SHOULD ONLY HAPPEN WHEN USING INFERENCE') - n_img = 1234 - self._raw_shape = (n_img, 3, resolution, resolution) - self.img_size = resolution - self._type = 'dir' - self._all_fnames = [None for i in range(n_img)] - self._image_fnames = self._all_fnames - name = os.path.splitext(os.path.basename(path))[0] - print( - '==> use image path: %s, num images: %d' % ( - self.root, len(self._all_fnames))) - super().__init__(name=name, raw_shape=self._raw_shape, **super_kwargs) - return - folder_list = sorted(os.listdir(root)) - if data_camera_mode == 'shapenet_chair' or data_camera_mode == 'shapenet_car': - if data_camera_mode == 'shapenet_car': - split_name = './3dgan_data_split/shapenet_car/%s.txt' % (split) - if split == 'all': - split_name = './3dgan_data_split/shapenet_car.txt' - elif data_camera_mode == 'shapenet_chair': - split_name = './3dgan_data_split/shapenet_chair/%s.txt' % (split) - if split == 'all': - split_name = './3dgan_data_split/shapenet_chair.txt' - valid_folder_list = [] - with open(split_name, 'r') as f: - all_line = f.readlines() - for l in all_line: - valid_folder_list.append(l.strip()) - valid_folder_list = set(valid_folder_list) - useful_folder_list = set(folder_list).intersection(valid_folder_list) - folder_list = sorted(list(useful_folder_list)) - if data_camera_mode == 'ts_animal': - split_name = './3dgan_data_split/ts_animals/%s.txt' % (split) - print('==> use ts animal split %s' % (split)) - if split != 'all': - valid_folder_list = [] - with open(split_name, 'r') as f: - all_line = f.readlines() - for l in all_line: - valid_folder_list.append(l.strip()) - valid_folder_list = set(valid_folder_list) - useful_folder_list = set(folder_list).intersection(valid_folder_list) - folder_list = sorted(list(useful_folder_list)) - elif data_camera_mode == 'shapenet_motorbike': - split_name = './3dgan_data_split/shapenet_motorbike/%s.txt' % (split) - print('==> use ts shapenet motorbike split %s' % (split)) - if split != 'all': - valid_folder_list = [] - with open(split_name, 'r') as f: - all_line = f.readlines() - for l in all_line: - valid_folder_list.append(l.strip()) - valid_folder_list = set(valid_folder_list) - useful_folder_list = set(folder_list).intersection(valid_folder_list) - folder_list = sorted(list(useful_folder_list)) + + if not os.path.exists(root): + print('==> WARNING: Root path does not exist. If you are running inference this is OK: %s' % root) + n_img = 1234 + self._raw_shape = (n_img, 3, resolution, resolution) + self.img_size = resolution + self._type = 'dir' + self._all_fnames = [None for i in range(n_img)] + self._image_fnames = self._all_fnames + name = os.path.splitext(os.path.basename(path))[0] + print( + '==> use image path: %s, num images: %d' % ( + self.root, len(self._all_fnames))) + super().__init__(name=name, raw_shape=self._raw_shape, **super_kwargs) + return + folder_list = sorted(os.listdir(root)) + if manifest_dir.includes('shapenet'): + split_name = './3dgan_data_split/' + manifest_dir + '/%s.txt' % (split) + valid_folder_list = [] + with open(split_name, 'r') as f: + all_line = f.readlines() + for l in all_line: + valid_folder_list.append(l.strip()) + valid_folder_list = set(valid_folder_list) + useful_folder_list = set(folder_list).intersection(valid_folder_list) + folder_list = sorted(list(useful_folder_list)) print('==> use shapenet folder number %s' % (len(folder_list))) folder_list = [os.path.join(root, f) for f in folder_list] all_img_list = [] @@ -242,17 +207,15 @@ def __init__( self.img_list = all_img_list self.mask_list = all_mask_list - else: - raise NotImplementedError - self.img_size = resolution - self._type = 'dir' - self._all_fnames = self.img_list - self._image_fnames = self._all_fnames - name = os.path.splitext(os.path.basename(self._path))[0] - print( - '==> use image path: %s, num images: %d' % (self.root, len(self._all_fnames))) - raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape) - super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) + self.img_size = resolution + self._type = 'dir' + self._all_fnames = self.img_list + self._image_fnames = self._all_fnames + name = os.path.splitext(os.path.basename(self._path))[0] + print( + '==> use image path: %s, num images: %d' % (self.root, len(self._all_fnames))) + raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape) + super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) @staticmethod def _file_ext(fname): @@ -283,31 +246,22 @@ def __getstate__(self): def __getitem__(self, idx): fname = self._image_fnames[self._raw_idx[idx]] - if self.data_camera_mode == 'shapenet_car' or self.data_camera_mode == 'shapenet_chair' \ - or self.data_camera_mode == 'renderpeople' \ - or self.data_camera_mode == 'shapenet_motorbike' or self.data_camera_mode == 'ts_house' or self.data_camera_mode == 'ts_animal' \ - : - ori_img = cv2.imread(fname, cv2.IMREAD_UNCHANGED) - img = ori_img[:, :, :3][..., ::-1] - mask = ori_img[:, :, 3:4] - condinfo = np.zeros(2) - fname_list = fname.split('/') - img_idx = int(fname_list[-1].split('.')[0]) - obj_idx = fname_list[-2] - syn_idx = fname_list[-3] - - if self.data_camera_mode == 'shapenet_car' or self.data_camera_mode == 'shapenet_chair' \ - or self.data_camera_mode == 'renderpeople' or self.data_camera_mode == 'shapenet_motorbike' \ - or self.data_camera_mode == 'ts_house' or self.data_camera_mode == 'ts_animal': - if not os.path.exists(os.path.join(self.camera_root, syn_idx, obj_idx, 'rotation.npy')): - print('==> not found camera root') - else: - rotation_camera = np.load(os.path.join(self.camera_root, syn_idx, obj_idx, 'rotation.npy')) - elevation_camera = np.load(os.path.join(self.camera_root, syn_idx, obj_idx, 'elevation.npy')) - condinfo[0] = rotation_camera[img_idx] / 180 * np.pi - condinfo[1] = (90 - elevation_camera[img_idx]) / 180.0 * np.pi + ori_img = cv2.imread(fname, cv2.IMREAD_UNCHANGED) + img = ori_img[:, :, :3][..., ::-1] + mask = ori_img[:, :, 3:4] + condinfo = np.zeros(2) + fname_list = fname.split('/') + img_idx = int(fname_list[-1].split('.')[0]) + obj_idx = fname_list[-2] + syn_idx = fname_list[-3] + + if not os.path.exists(os.path.join(self.camera_root, syn_idx, obj_idx, 'rotation.npy')): + print('==> not found camera root') else: - raise NotImplementedError + rotation_camera = np.load(os.path.join(self.camera_root, syn_idx, obj_idx, 'rotation.npy')) + elevation_camera = np.load(os.path.join(self.camera_root, syn_idx, obj_idx, 'elevation.npy')) + condinfo[0] = rotation_camera[img_idx] / 180 * np.pi + condinfo[1] = (90 - elevation_camera[img_idx]) / 180.0 * np.pi resize_img = cv2.resize(img, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR) if not mask is None: diff --git a/training/discriminator_architecture.py b/training/discriminator_architecture.py index 241f4a9..5dd6a08 100644 --- a/training/discriminator_architecture.py +++ b/training/discriminator_architecture.py @@ -526,7 +526,6 @@ def __init__( conv_clamp=256, # Clamp the output of convolution layers to +-X, None = disable clamping. cmap_dim=None, # Dimensionality of mapped conditioning label, None = default. add_camera_cond=False, # whether add camera for conditioning - data_camera_mode='', # device='cuda', block_kwargs={}, # Arguments for DiscriminatorBlock. mapping_kwargs={}, # Arguments for MappingNetwork. @@ -534,7 +533,6 @@ def __init__( ): super().__init__() - self.data_camera_mode = data_camera_mode self.conditional_dim = c_dim self.c_dim = c_dim self.add_camera_cond = add_camera_cond @@ -549,15 +547,9 @@ def __init__( # We will compute the positional encoding for the camera if self.add_camera_cond: - if self.data_camera_mode == 'shapenet_car' \ - or self.data_camera_mode == 'shapenet_chair' or self.data_camera_mode == 'shapenet_motorbike' \ - or self.data_camera_mode == 'renderpeople' or self.data_camera_mode == 'ts_house' \ - or self.data_camera_mode == 'ts_animal': - self.camera_dim = 2 - self.camera_dim_enc = 2 + 2 * 2 * 3 + self.camera_dim = 2 + self.camera_dim_enc = 2 + 2 * 2 * 3 - else: - raise NotImplementedError self.c_dim = self.camera_dim_enc + self.c_dim if cmap_dim is None: @@ -609,19 +601,13 @@ def __init__( def pos_enc_angle(self, camera_angle): # Encode the camera angles into cos/sin - if self.data_camera_mode == 'shapenet_car' \ - or self.data_camera_mode == 'shapenet_chair' or self.data_camera_mode == 'shapenet_motorbike' \ - or self.data_camera_mode == 'renderpeople' or self.data_camera_mode == 'ts_house' \ - or self.data_camera_mode == 'ts_animal': - L = 3 - p_transformed = torch.cat( - [torch.cat( - [torch.sin((2 ** i) * camera_angle), - torch.cos((2 ** i) * camera_angle)], - dim=-1) for i in range(L)], dim=-1) - p_transformed = torch.cat([p_transformed, camera_angle], dim=-1) - else: - raise NotImplementedError + L = 3 + p_transformed = torch.cat( + [torch.cat( + [torch.sin((2 ** i) * camera_angle), + torch.cos((2 ** i) * camera_angle)], + dim=-1) for i in range(L)], dim=-1) + p_transformed = torch.cat([p_transformed, camera_angle], dim=-1) return p_transformed def forward(self, img, c, update_emas=False, alpha=1.0, mask_pyramid=None, **block_kwargs): diff --git a/training/loss.py b/training/loss.py index b0f16a8..9931fb6 100644 --- a/training/loss.py +++ b/training/loss.py @@ -105,15 +105,7 @@ def accumulate_gradients( gen_z, gen_c, return_shape=True ) - camera_condition = None - if self.G.synthesis.data_camera_mode == 'shapenet_car' or self.G.synthesis.data_camera_mode == 'shapenet_chair' \ - or self.G.synthesis.data_camera_mode == 'shapenet_motorbike' or self.G.synthesis.data_camera_mode == 'renderpeople' or \ - self.G.synthesis.data_camera_mode == 'shapenet_plant' or self.G.synthesis.data_camera_mode == 'shapenet_vase' or \ - self.G.synthesis.data_camera_mode == 'ts_house' or self.G.synthesis.data_camera_mode == 'ts_animal' or \ - self.G.synthesis.data_camera_mode == 'all_shapenet': - camera_condition = torch.cat((gen_camera[-2], gen_camera[-1]), dim=-1) - else: - assert NotImplementedError + camera_condition = torch.cat((gen_camera[-2], gen_camera[-1]), dim=-1) # Send to discriminator gen_logits = self.run_D(gen_img, camera_condition, mask_pyramid=mask_pyramid) gen_logits, gen_logits_mask = gen_logits @@ -150,14 +142,8 @@ def accumulate_gradients( # First generate the rendered image of generated 3D shapes gen_img, _gen_ws, gen_camera, mask_pyramid, render_return_value = self.run_G( gen_z, gen_c, update_emas=True) - if self.G.synthesis.data_camera_mode == 'shapenet_car' or self.G.synthesis.data_camera_mode == 'shapenet_chair' \ - or self.G.synthesis.data_camera_mode == 'shapenet_motorbike' or self.G.synthesis.data_camera_mode == 'renderpeople' or \ - self.G.synthesis.data_camera_mode == 'shapenet_plant' or self.G.synthesis.data_camera_mode == 'shapenet_vase' or \ - self.G.synthesis.data_camera_mode == 'ts_house' or self.G.synthesis.data_camera_mode == 'ts_animal' or \ - self.G.synthesis.data_camera_mode == 'all_shapenet': - camera_condition = torch.cat((gen_camera[-2], gen_camera[-1]), dim=-1) - else: - camera_condition = None + + camera_condition = torch.cat((gen_camera[-2], gen_camera[-1]), dim=-1) # Send it to discriminator gen_logits = self.run_D( diff --git a/training/networks_get3d.py b/training/networks_get3d.py index 67175ce..1965a90 100644 --- a/training/networks_get3d.py +++ b/training/networks_get3d.py @@ -28,7 +28,6 @@ def __init__( img_resolution, # Output image resolution. img_channels, # Number of color channels. device='cuda', - data_camera_mode='carla', geometry_type='normal', tet_res=64, # Resolution for tetrahedron grid render_type='neural_render', # neural type @@ -52,7 +51,6 @@ def __init__( self.deformation_multiplier = deformation_multiplier self.geometry_type = geometry_type - self.data_camera_mode = data_camera_mode self.n_freq_posenc_geo = 1 self.render_type = render_type @@ -318,7 +316,7 @@ def generate_random_camera(self, batch_size, n_views=2): ''' sample_r = None world2cam_matrix, forward_vector, camera_origin, rotation_angle, elevation_angle = sample_camera( - self.data_camera_mode, batch_size * n_views, self.device) + batch_size * n_views, self.device) mv_batch = world2cam_matrix campos = camera_origin return campos.reshape(batch_size, n_views, 3), mv_batch.reshape(batch_size, n_views, 4, 4), \ diff --git a/training/sample_camera_distribution.py b/training/sample_camera_distribution.py index 868d88b..74c063b 100644 --- a/training/sample_camera_distribution.py +++ b/training/sample_camera_distribution.py @@ -32,38 +32,17 @@ def create_camera_from_angle(phi, theta, sample_r, device='cuda'): return world2cam_matrix, forward_vector, camera_origin, phi, theta -def sample_camera(camera_data_mode, n, device='cuda'): +def sample_camera(n, device='cuda'): # We use this function to sample the camera for training the generator # When we're rendering the dataset, the camera is random sampled from # a uniform distribution (see `render_shapenet_data/render_shapenet.py` # for details of camera distribution) - if camera_data_mode == 'shapenet_car' or camera_data_mode == 'shapenet_chair' or camera_data_mode == 'shapenet_motorbike' \ - or camera_data_mode == 'ts_house': - horizontal_stddev = math.pi # here means horizontal rotation - vertical_stddev = (math.pi / 180) * 15 - horizontal_mean = math.pi ######## [horizon range [0, 2pi]] - vertical_mean = (math.pi / 180) * 75 - mode = 'uniform' - radius_range = [1.2, 1.2] - elif camera_data_mode == 'ts_animal': - horizontal_stddev = math.pi # here means horizontal rotation - vertical_stddev = (math.pi / 180) * (45.0 / 2.0) - - horizontal_mean = math.pi ######## [horizon range [0, 2pi]] - vertical_mean = (math.pi / 180) * ((90.0 + 90.0 - 45.0) / 2.0) - mode = 'uniform' - radius_range = [1.2, 1.2] - elif camera_data_mode == 'renderpeople': - horizontal_stddev = math.pi # here means horizontal rotation - vertical_stddev = (math.pi / 180) * 10 - - horizontal_mean = math.pi ######## [horizon range [0, 2pi]] - vertical_mean = (math.pi / 180) * 80 - mode = 'uniform' - radius_range = [1.2, 1.2] - - else: - raise NotImplementedError + horizontal_stddev = math.pi # here means horizontal rotation + vertical_stddev = (math.pi / 180) * 15 + horizontal_mean = math.pi ######## [horizon range [0, 2pi]] + vertical_mean = (math.pi / 180) * 80 + mode = 'uniform' + radius_range = [1.2, 1.2] camera_origin, rotation_angle, elevation_angle = sample_camera_positions( device, n=n, r=radius_range, From 43cc2bc2c7631010b5fae9a342a856c8df41354f Mon Sep 17 00:00:00 2001 From: shawticus Date: Mon, 10 Oct 2022 05:34:25 -0700 Subject: [PATCH 2/5] Remove prints --- train_3d.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/train_3d.py b/train_3d.py index 55dede0..852a895 100644 --- a/train_3d.py +++ b/train_3d.py @@ -22,13 +22,10 @@ # ---------------------------------------------------------------------------- def subprocess_fn(rank, c, temp_dir): - print('Subprocess...1') dnnlib.util.Logger(file_name=os.path.join(c.run_dir, 'log.txt'), file_mode='a', should_flush=True) - print('Subprocess...2') # Init torch.distributed. if c.num_gpus > 1: - print('Subprocess...3') init_file = os.path.abspath(os.path.join(temp_dir, '.torch_distributed_init')) if os.name == 'nt': init_method = 'file:///' + init_file.replace('\\', '/') @@ -38,10 +35,8 @@ def subprocess_fn(rank, c, temp_dir): init_method = f'file://{init_file}' torch.distributed.init_process_group( backend='nccl', init_method=init_method, rank=rank, world_size=c.num_gpus) - print('Subprocess...4') # Init torch_utils. sync_device = torch.device('cuda', rank) if c.num_gpus > 1 else None - print('Subprocess...5') training_stats.init_multiprocessing(rank=rank, sync_device=sync_device) if rank != 0: custom_ops.verbosity = 'none' @@ -100,15 +95,11 @@ def launch_training(c, desc, outdir, dry_run): json.dump(c, f, indent=2) # Launch processes. - print('Launching processes...') torch.multiprocessing.set_start_method('spawn', force=True) - print('Launching processes...2') - with tempfile.TemporaryDirectory() as temp_dir: if c.num_gpus == 1: subprocess_fn(rank=0, c=c, temp_dir=temp_dir) else: - print('Launching processes...3') torch.multiprocessing.spawn(fn=subprocess_fn, args=(c, temp_dir), nprocs=c.num_gpus) From ceb97c7f15a24cacb75570510e02e7395e2e2342 Mon Sep 17 00:00:00 2001 From: shawticus Date: Mon, 10 Oct 2022 05:44:36 -0700 Subject: [PATCH 3/5] Fix indentation and formatting --- train_3d.py | 2 ++ training/dataset.py | 18 +++++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/train_3d.py b/train_3d.py index 852a895..075279f 100644 --- a/train_3d.py +++ b/train_3d.py @@ -35,6 +35,7 @@ def subprocess_fn(rank, c, temp_dir): init_method = f'file://{init_file}' torch.distributed.init_process_group( backend='nccl', init_method=init_method, rank=rank, world_size=c.num_gpus) + # Init torch_utils. sync_device = torch.device('cuda', rank) if c.num_gpus > 1 else None training_stats.init_multiprocessing(rank=rank, sync_device=sync_device) @@ -95,6 +96,7 @@ def launch_training(c, desc, outdir, dry_run): json.dump(c, f, indent=2) # Launch processes. + print('Launching processes...') torch.multiprocessing.set_start_method('spawn', force=True) with tempfile.TemporaryDirectory() as temp_dir: if c.num_gpus == 1: diff --git a/training/dataset.py b/training/dataset.py index d47b14a..d203728 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -207,15 +207,15 @@ def __init__( self.img_list = all_img_list self.mask_list = all_mask_list - self.img_size = resolution - self._type = 'dir' - self._all_fnames = self.img_list - self._image_fnames = self._all_fnames - name = os.path.splitext(os.path.basename(self._path))[0] - print( - '==> use image path: %s, num images: %d' % (self.root, len(self._all_fnames))) - raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape) - super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) + self.img_size = resolution + self._type = 'dir' + self._all_fnames = self.img_list + self._image_fnames = self._all_fnames + name = os.path.splitext(os.path.basename(self._path))[0] + print( + '==> use image path: %s, num images: %d' % (self.root, len(self._all_fnames))) + raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape) + super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) @staticmethod def _file_ext(fname): From ae8995cbea2d97293f4218266a6f55f55bedef2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=CC=B5=CD=9D=CC=85=CC=8F=CC=8E=CC=9E=CC=97=CC=9D=CC=BCO?= =?UTF-8?q?=CC=B4=CC=87=CC=8A=CC=83=CC=8B=CC=80=CC=9D=CC=BBO=CC=B7=CC=83?= =?UTF-8?q?=CD=8B=CC=BCN=CC=B8=CC=BF=CD=9C=CC=A9=20=CC=B6=CD=92=CC=9C?= =?UTF-8?q?=CC=A0=CC=B9=CC=BC=CC=A9?= Date: Tue, 11 Oct 2022 04:27:05 -0700 Subject: [PATCH 4/5] Update train_3d.py --- train_3d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train_3d.py b/train_3d.py index 075279f..12334e0 100644 --- a/train_3d.py +++ b/train_3d.py @@ -249,7 +249,7 @@ def main(**kwargs): c.G_kwargs.render_type = opts.render_type c.G_kwargs.use_tri_plane = opts.use_tri_plane - c.D_kwargs.manifest_dir = opts.manifest_dir + # c.D_kwargs.manifest_dir = opts.manifest_dir c.D_kwargs.add_camera_cond = opts.add_camera_cond c.G_kwargs.tet_res = opts.tet_res @@ -259,7 +259,7 @@ def main(**kwargs): c.batch_size = opts.batch c.batch_gpu = opts.batch_gpu or opts.batch // opts.gpus # c.G_kwargs.geo_pos_enc = opts.geo_pos_enc - c.G_kwargs.manifest_dir = opts.manifest_dir + # c.G_kwargs.manifest_dir = opts.manifest_dir c.G_kwargs.channel_base = c.D_kwargs.channel_base = opts.cbase c.G_kwargs.channel_max = c.D_kwargs.channel_max = opts.cmax From 10757f3ff51fffe580f5f0dc349749f1a71f055b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=CC=B5=CD=9D=CC=85=CC=8F=CC=8E=CC=9E=CC=97=CC=9D=CC=BCO?= =?UTF-8?q?=CC=B4=CC=87=CC=8A=CC=83=CC=8B=CC=80=CC=9D=CC=BBO=CC=B7=CC=83?= =?UTF-8?q?=CD=8B=CC=BCN=CC=B8=CC=BF=CD=9C=CC=A9=20=CC=B6=CD=92=CC=9C?= =?UTF-8?q?=CC=A0=CC=B9=CC=BC=CC=A9?= Date: Tue, 11 Oct 2022 04:28:30 -0700 Subject: [PATCH 5/5] Update dataset.py --- training/dataset.py | 45 ++++++++++++++++++++++----------------------- 1 file changed, 22 insertions(+), 23 deletions(-) diff --git a/training/dataset.py b/training/dataset.py index d203728..883236d 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -183,29 +183,28 @@ def __init__( super().__init__(name=name, raw_shape=self._raw_shape, **super_kwargs) return folder_list = sorted(os.listdir(root)) - if manifest_dir.includes('shapenet'): - split_name = './3dgan_data_split/' + manifest_dir + '/%s.txt' % (split) - valid_folder_list = [] - with open(split_name, 'r') as f: - all_line = f.readlines() - for l in all_line: - valid_folder_list.append(l.strip()) - valid_folder_list = set(valid_folder_list) - useful_folder_list = set(folder_list).intersection(valid_folder_list) - folder_list = sorted(list(useful_folder_list)) - print('==> use shapenet folder number %s' % (len(folder_list))) - folder_list = [os.path.join(root, f) for f in folder_list] - all_img_list = [] - all_mask_list = [] - - for folder in folder_list: - rgb_list = sorted(os.listdir(folder)) - rgb_file_name_list = [os.path.join(folder, n) for n in rgb_list] - all_img_list.extend(rgb_file_name_list) - all_mask_list.extend(rgb_list) - - self.img_list = all_img_list - self.mask_list = all_mask_list + split_name = './3dgan_data_split/' + manifest_dir + '/%s.txt' % (split) + valid_folder_list = [] + with open(split_name, 'r') as f: + all_line = f.readlines() + for l in all_line: + valid_folder_list.append(l.strip()) + valid_folder_list = set(valid_folder_list) + useful_folder_list = set(folder_list).intersection(valid_folder_list) + folder_list = sorted(list(useful_folder_list)) + print('==> use shapenet folder number %s' % (len(folder_list))) + folder_list = [os.path.join(root, f) for f in folder_list] + all_img_list = [] + all_mask_list = [] + + for folder in folder_list: + rgb_list = sorted(os.listdir(folder)) + rgb_file_name_list = [os.path.join(folder, n) for n in rgb_list] + all_img_list.extend(rgb_file_name_list) + all_mask_list.extend(rgb_list) + + self.img_list = all_img_list + self.mask_list = all_mask_list self.img_size = resolution self._type = 'dir'