diff --git a/convert.py b/convert.py index 78948848f..9351c6458 100644 --- a/convert.py +++ b/convert.py @@ -8,12 +8,16 @@ # # For inquiries contact george.drettakis@inria.fr # +# xvdp removed magick, even single threaded PIL resizes 4X faster + import os import logging from argparse import ArgumentParser import shutil +from PIL import Image + # This Python script is based on the shell converter script provided in the MipNerF 360 repository. parser = ArgumentParser("Colmap converter") parser.add_argument("--no_gpu", action='store_true') @@ -25,7 +29,7 @@ parser.add_argument("--magick_executable", default="", type=str) args = parser.parse_args() colmap_command = '"{}"'.format(args.colmap_executable) if len(args.colmap_executable) > 0 else "colmap" -magick_command = '"{}"'.format(args.magick_executable) if len(args.magick_executable) > 0 else "magick" + use_gpu = 1 if not args.no_gpu else 0 if not args.skip_matching: @@ -87,38 +91,21 @@ destination_file = os.path.join(args.source_path, "sparse", "0", file) shutil.move(source_file, destination_file) -if(args.resize): +if args.resize: print("Copying and resizing...") # Resize images. - os.makedirs(args.source_path + "/images_2", exist_ok=True) - os.makedirs(args.source_path + "/images_4", exist_ok=True) - os.makedirs(args.source_path + "/images_8", exist_ok=True) + for div in [2,4,8]: + os.makedirs(args.source_path + f"/images_{div}", exist_ok=True) # Get the list of files in the source directory files = os.listdir(args.source_path + "/images") # Copy each file from the source directory to the destination directory - for file in files: + for j, file in enumerate(files): source_file = os.path.join(args.source_path, "images", file) - - destination_file = os.path.join(args.source_path, "images_2", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system(magick_command + " mogrify -resize 50% " + destination_file) - if exit_code != 0: - logging.error(f"50% resize failed with code {exit_code}. Exiting.") - exit(exit_code) - - destination_file = os.path.join(args.source_path, "images_4", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system(magick_command + " mogrify -resize 25% " + destination_file) - if exit_code != 0: - logging.error(f"25% resize failed with code {exit_code}. Exiting.") - exit(exit_code) - - destination_file = os.path.join(args.source_path, "images_8", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system(magick_command + " mogrify -resize 12.5% " + destination_file) - if exit_code != 0: - logging.error(f"12.5% resize failed with code {exit_code}. Exiting.") - exit(exit_code) + im = Image.open(source_file) + logging.info(f"processing image [{j}/{len(files)}] {source_file}") + for div in [2,4,8]: + destination_file = os.path.join(args.source_path, f"images_{div}", file) + im.resize([round(i/div) for i in im.size], Image.BICUBIC).save(destination_file, quality=100) print("Done.") diff --git a/scene/cameras.py b/scene/cameras.py index abf6e5242..bc59fcc38 100644 --- a/scene/cameras.py +++ b/scene/cameras.py @@ -56,6 +56,16 @@ def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) self.camera_center = self.world_view_transform.inverse()[3, :3] + def __repr__(self): + format_string = self.__class__.__name__ + '()\n' + for k, v in self.__dict__.items(): + if torch.is_tensor(v) and v.numel() > 16: + format_string +=f" {k}:\t{tuple(v.shape)}\n" + else: + format_string += f"{k}:\t{v}\n" + return format_string + + class MiniCam: def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): self.image_width = width diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index 632a1e8e1..07a30d0bf 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -58,6 +58,17 @@ def __init__(self, sh_degree : int): self.spatial_lr_scale = 0 self.setup_functions() + + def __repr__(self): + format_string = self.__class__.__name__ + '()\n' + for k, v in self.__dict__.items(): + if torch.is_tensor(v): + format_string +=f" {k}:\t{tuple(v.shape)}\n" + else: + format_string += f"{k}:\t{v}\n" + return format_string + + def capture(self): return ( self.active_sh_degree, @@ -73,7 +84,7 @@ def capture(self): self.optimizer.state_dict(), self.spatial_lr_scale, ) - + def restore(self, model_args, training_args): (self.active_sh_degree, self._xyz, diff --git a/train.py b/train.py index 5d819b348..d57acae68 100644 --- a/train.py +++ b/train.py @@ -48,7 +48,7 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi ema_loss_for_log = 0.0 progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress") first_iter += 1 - for iteration in range(first_iter, opt.iterations + 1): + for iteration in range(first_iter, opt.iterations + 1): if network_gui.conn == None: network_gui.try_connect() while network_gui.conn != None: @@ -62,7 +62,10 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi if do_training and ((iteration < int(opt.iterations)) or not keep_alive): break except Exception as e: + network_gui.conn.close() network_gui.conn = None + network_gui.listener.close() + network_gui.listener = None iter_start.record() @@ -162,7 +165,7 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i # Report test and samples of training set if iteration in testing_iterations: torch.cuda.empty_cache() - validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()}, + validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()}, {'name': 'train', 'cameras' : [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in range(5, 30, 5)]}) for config in validation_configs: @@ -203,11 +206,12 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i parser.add_argument("--test_iterations", nargs="+", type=int, default=[7_000, 30_000]) parser.add_argument("--save_iterations", nargs="+", type=int, default=[7_000, 30_000]) parser.add_argument("--quiet", action="store_true") - parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[]) + parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[7_000, 30_000]) parser.add_argument("--start_checkpoint", type=str, default = None) + args = parser.parse_args(sys.argv[1:]) args.save_iterations.append(args.iterations) - + print("Optimizing " + args.model_path) # Initialize system state (RNG) @@ -216,7 +220,8 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i # Start GUI server, configure and run training network_gui.init(args.ip, args.port) torch.autograd.set_detect_anomaly(args.detect_anomaly) - training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from) + training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, + args.checkpoint_iterations, args.start_checkpoint, args.debug_from) # All done print("\nTraining complete.") diff --git a/utils/general_utils.py b/utils/general_utils.py index 541c08252..43f56eee7 100644 --- a/utils/general_utils.py +++ b/utils/general_utils.py @@ -18,13 +18,15 @@ def inverse_sigmoid(x): return torch.log(x/(1-x)) -def PILtoTorch(pil_image, resolution): +def PILtoTorch(pil_image, resolution, pin_memory=True): resized_image_PIL = pil_image.resize(resolution) - resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 - if len(resized_image.shape) == 3: - return resized_image.permute(2, 0, 1) - else: - return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) + resized_image = torch.from_numpy(np.array(resized_image_PIL, dtype=np.float32)) / 255.0 + if resized_image.ndim == 2: + resized_image = resized_image[None] + resized_image = resized_image.permute(2, 0, 1).contiguous() + if pin_memory: + resized_image.pin_memory = True + return resized_image def get_expon_lr_func( lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000