From e35781fe858113caa24675a239a14d1336a4da4a Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Sun, 31 Jul 2022 18:22:57 -0700 Subject: [PATCH] Code cleanup (#234) --- nerfactory/data/colmap_utils.py | 16 +++---- nerfactory/data/pixel_sampler.py | 2 +- nerfactory/fields/base.py | 6 +-- nerfactory/fields/modules/encoding.py | 2 +- nerfactory/graphs/modules/ray_sampler.py | 2 +- nerfactory/utils/io.py | 8 ++-- nerfactory/viewer/server/geometry.py | 57 +++++++++--------------- nerfactory/viewer/server/path.py | 4 +- nerfactory/viewer/server/server.py | 1 - nerfactory/viewer/server/tree.py | 6 +-- scripts/run_eval.py | 2 +- scripts/run_train.py | 2 +- tests/utils/test_tensor_dataclass.py | 2 +- tests/viewer/test_viewer.py | 4 +- 14 files changed, 46 insertions(+), 68 deletions(-) diff --git a/nerfactory/data/colmap_utils.py b/nerfactory/data/colmap_utils.py index deb3b510e..6cc8b2e6d 100644 --- a/nerfactory/data/colmap_utils.py +++ b/nerfactory/data/colmap_utils.py @@ -66,8 +66,8 @@ CAMERA_MODELS = { CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), } -CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]) -CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]) +CAMERA_MODEL_IDS = {camera_model.model_id: camera_model for camera_model in CAMERA_MODELS} +CAMERA_MODEL_NAMES = {camera_model.model_name: camera_model for camera_model in CAMERA_MODELS} def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): @@ -109,7 +109,7 @@ def read_cameras_text(path): void Reconstruction::ReadCamerasText(const std::string& path) """ cameras = {} - with open(path, "r") as fid: + with open(path) as fid: while True: line = fid.readline() if not line: @@ -160,7 +160,7 @@ def write_cameras_text(cameras, path): HEADER = ( "# Camera list with one line of data per camera:\n" + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" - + "# Number of cameras: {}\n".format(len(cameras)) + + f"# Number of cameras: {len(cameras)}\n" ) with open(path, "w") as fid: fid.write(HEADER) @@ -194,7 +194,7 @@ def read_images_text(path): void Reconstruction::WriteImagesText(const std::string& path) """ images = {} - with open(path, "r") as fid: + with open(path) as fid: while True: line = fid.readline() if not line: @@ -272,7 +272,7 @@ def write_images_text(images, path): "# Image list with two lines of data per image:\n" + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" - + "# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations) + + f"# Number of images: {len(images)}, mean observations per image: {mean_observations}\n" ) with open(path, "w") as fid: @@ -316,7 +316,7 @@ def read_pointsTD_text(path): void Reconstruction::WritePoints3DText(const std::string& path) """ points3D = {} - with open(path, "r") as fid: + with open(path) as fid: while True: line = fid.readline() if not line: @@ -374,7 +374,7 @@ def write_pointsTD_text(points3D, path): HEADER = ( "# 3D point list with one line of data per point:\n" + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" - + "# Number of points: {}, mean track length: {}\n".format(len(points3D), mean_track_length) + + f"# Number of points: {len(points3D)}, mean track length: {mean_track_length}\n" ) with open(path, "w") as fid: diff --git a/nerfactory/data/pixel_sampler.py b/nerfactory/data/pixel_sampler.py index 9b7390b0b..61bd1df95 100644 --- a/nerfactory/data/pixel_sampler.py +++ b/nerfactory/data/pixel_sampler.py @@ -41,7 +41,7 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool * torch.tensor([num_images, image_height, image_width], device=device) ).long() - c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] + c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) image = batch["image"][c, y, x] mask, semantics = None, None if "mask" in batch: diff --git a/nerfactory/fields/base.py b/nerfactory/fields/base.py index 021d49f00..802a34855 100644 --- a/nerfactory/fields/base.py +++ b/nerfactory/fields/base.py @@ -37,9 +37,9 @@ class Field(nn.Module): frustums=Frustums( origins=positions, directions=torch.ones_like(positions), - starts=torch.zeros_like((positions[..., :1])), - ends=torch.zeros_like((positions[..., :1])), - pixel_area=torch.ones_like((positions[..., :1])), + starts=torch.zeros_like(positions[..., :1]), + ends=torch.zeros_like(positions[..., :1]), + pixel_area=torch.ones_like(positions[..., :1]), ) ) density, _ = self.get_density(ray_samples) diff --git a/nerfactory/fields/modules/encoding.py b/nerfactory/fields/modules/encoding.py index 914ea7b34..c195b9f75 100644 --- a/nerfactory/fields/modules/encoding.py +++ b/nerfactory/fields/modules/encoding.py @@ -192,7 +192,7 @@ class RFFEncoding(Encoding): if covs is None: encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1)) else: - input_var = torch.sum(((covs @ self.b_matrix)) * self.b_matrix, -2) + input_var = torch.sum((covs @ self.b_matrix) * self.b_matrix, -2) encoded_inputs = expected_sin( torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1) ) diff --git a/nerfactory/graphs/modules/ray_sampler.py b/nerfactory/graphs/modules/ray_sampler.py index 313d16755..e4cf53dc9 100644 --- a/nerfactory/graphs/modules/ray_sampler.py +++ b/nerfactory/graphs/modules/ray_sampler.py @@ -105,7 +105,7 @@ class SpacedSampler(Sampler): bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...] # [1, num_samples+1] - s_near, s_far = [self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars)] + s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars)) bins = self.spacing_fn_inv(bins * s_far + (1 - bins) * s_near) # [num_rays, num_samples+1] if self.train_stratified and self.training: diff --git a/nerfactory/utils/io.py b/nerfactory/utils/io.py index 7a0e89fd7..62f7d8719 100644 --- a/nerfactory/utils/io.py +++ b/nerfactory/utils/io.py @@ -29,7 +29,7 @@ def load_from_json(filename: str): filename: The filename to load from. """ assert filename.endswith(".json") - with open(filename, "r", encoding="UTF-8") as file: + with open(filename, encoding="UTF-8") as file: return json.load(file) @@ -102,10 +102,8 @@ def get_project_root(path: str) -> str: project_root = os.getenv("PROJECT_ROOT") if project_root is None: logging.info( - ( - "Please set PROJECT_ROOT to the root directory of this repo. " - "Going to try calling get_git_root(path) instead." - ) + "Please set PROJECT_ROOT to the root directory of this repo. " + "Going to try calling get_git_root(path) instead." ) project_root = get_git_root(path) return project_root diff --git a/nerfactory/viewer/server/geometry.py b/nerfactory/viewer/server/geometry.py index 8678fe867..a0a822f74 100644 --- a/nerfactory/viewer/server/geometry.py +++ b/nerfactory/viewer/server/geometry.py @@ -15,7 +15,6 @@ """Geometry to render meshes, textures, etc. """ -from __future__ import absolute_import, division, print_function import base64 import sys @@ -25,13 +24,8 @@ from typing import Any, BinaryIO, Dict, Iterable, List, Optional, Tuple import numpy as np import umsgpack -if sys.version_info >= (3, 0): - UNICODE = str - from io import BytesIO, StringIO -else: - from StringIO import StringIO - - BytesIO = StringIO +UNICODE = str +from io import BytesIO, StringIO class SceneElement: @@ -434,7 +428,7 @@ class OrthographicCamera(SceneElement): """ def __init__(self, left: float, right: float, top: float, bottom: float, near: float, far: float, zoom: int = 1): - super(OrthographicCamera, self).__init__() + super().__init__() self.left = left self.right = right self.top = top @@ -492,7 +486,7 @@ class PerspectiveCamera(SceneElement): film_offset: float = 0, focus: float = 10, ): - super(PerspectiveCamera, self).__init__() + super().__init__() self.fov = fov self.aspect = aspect self.far = far @@ -583,15 +577,12 @@ def data_from_stream(stream: BinaryIO) -> str: Raises: ValueError: invalid stream type """ - if sys.version_info >= (3, 0): - if isinstance(stream, BytesIO): - data = stream.read().decode(encoding="utf-8") - elif isinstance(stream, StringIO): - data = stream.read() - else: - raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(stream)}") - else: + if isinstance(stream, BytesIO): + data = stream.read().decode(encoding="utf-8") + elif isinstance(stream, StringIO): data = stream.read() + else: + raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(stream)}") return data @@ -629,7 +620,7 @@ class ObjMeshGeometry(MeshGeometry): Args: fname: file name to read from """ - with open(fname, "r", encoding="utf8") as f: + with open(fname, encoding="utf8") as f: return MeshGeometry(f.read(), "obj") @staticmethod @@ -659,7 +650,7 @@ class DaeMeshGeometry(MeshGeometry): Args: fname: file name to read from """ - with open(fname, "r", encoding="utf8") as f: + with open(fname, encoding="utf8") as f: return MeshGeometry(f.read(), "dae") @staticmethod @@ -702,15 +693,12 @@ class StlMeshGeometry(MeshGeometry): Args: f: binary io stream to read from """ - if sys.version_info >= (3, 0): - if isinstance(f, BytesIO): - arr = np.frombuffer(f.read(), dtype=np.uint8) - elif isinstance(f, StringIO): - arr = np.frombuffer(bytes(f.read(), "utf-8"), dtype=np.uint8) - else: - raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(f)}") - else: + if isinstance(f, BytesIO): arr = np.frombuffer(f.read(), dtype=np.uint8) + elif isinstance(f, StringIO): + arr = np.frombuffer(bytes(f.read(), "utf-8"), dtype=np.uint8) + else: + raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(f)}") _, extcode = threejs_type(np.uint8) encoded = umsgpack.Ext(extcode, arr.tobytes()) return MeshGeometry(encoded, "stl") @@ -746,15 +734,12 @@ class PlyMeshGeometry(MeshGeometry): Args: f: binary io stream to read from """ - if sys.version_info >= (3, 0): - if isinstance(f, BytesIO): - arr = np.frombuffer(f.read(), dtype=np.uint8) - elif isinstance(f, StringIO): - arr = np.frombuffer(bytes(f.read(), "utf-8"), dtype=np.uint8) - else: - raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(f)}") - else: + if isinstance(f, BytesIO): arr = np.frombuffer(f.read(), dtype=np.uint8) + elif isinstance(f, StringIO): + arr = np.frombuffer(bytes(f.read(), "utf-8"), dtype=np.uint8) + else: + raise ValueError(f"Stream must be instance of StringIO or BytesIO, not {type(f)}") _, extcode = threejs_type(np.uint8) encoded = umsgpack.Ext(extcode, arr.tobytes()) return MeshGeometry(encoded, "stl") diff --git a/nerfactory/viewer/server/path.py b/nerfactory/viewer/server/path.py index a74c7300f..ca80c5d29 100644 --- a/nerfactory/viewer/server/path.py +++ b/nerfactory/viewer/server/path.py @@ -15,13 +15,11 @@ """Path class """ -from __future__ import absolute_import, division, print_function import sys from typing import Tuple -if sys.version_info >= (3, 0): - UNICODE = str +UNICODE = str class Path: diff --git a/nerfactory/viewer/server/server.py b/nerfactory/viewer/server/server.py index 8201bf9d5..e37ccba24 100644 --- a/nerfactory/viewer/server/server.py +++ b/nerfactory/viewer/server/server.py @@ -14,7 +14,6 @@ """Server bridge to faciliate interactions between python backend and javascript front end""" -from __future__ import absolute_import, division, print_function import sys from typing import Callable, List, Optional diff --git a/nerfactory/viewer/server/tree.py b/nerfactory/viewer/server/tree.py index e18654fda..1eacaa766 100644 --- a/nerfactory/viewer/server/tree.py +++ b/nerfactory/viewer/server/tree.py @@ -16,7 +16,6 @@ Functions for traversing server bridge tree nodes (storage) """ -from __future__ import absolute_import, division, print_function from collections import defaultdict from typing import List @@ -28,7 +27,7 @@ class TreeNode(defaultdict): __slots__ = ["object", "transform", "properties", "animation"] def __init__(self, *args, **kwargs): - super(TreeNode, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) self.object = None self.properties = [] self.transform = None @@ -46,8 +45,7 @@ def walk(tree: TreeNode) -> None: """ yield tree for v in tree.values(): - for t in walk(v): # could use `yield from` if we didn't need python2 - yield t + yield from walk(v) def find_node(tree: TreeNode, path: List[str]) -> TreeNode: diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 4fc4dd3a7..1d3941efe 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -52,7 +52,7 @@ def _load_checkpoint(config: DictConfig, graph: Graph) -> None: if config.load_step is None: print("Loading latest checkpoint from load_dir") # NOTE: this is specific to the checkpoint name format - load_step = sorted([int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir)])[-1] + load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1] else: load_step = config.load_step load_path = os.path.join(config.load_dir, f"step-{load_step:09d}.ckpt") diff --git a/scripts/run_train.py b/scripts/run_train.py index 9d449a926..3e3dfb92c 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -212,7 +212,7 @@ def main(config: DictConfig): unrolled_path = os.path.join(os.getcwd(), ".hydra/config.yaml") if os.path.exists(unrolled_path): - with open(unrolled_path, "r", encoding="utf8") as f: + with open(unrolled_path, encoding="utf8") as f: unrolled_config = yaml.safe_load(f) logger = logging.getLogger(__name__) logger.info("Printing current config setup") diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 0836b52af..2c79f4992 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/tests/utils/test_tensor_dataclass.py @@ -48,7 +48,7 @@ def test_broadcasting(): assert tensor_dataclass.b.shape == (4, 6, 2) a = torch.ones((4, 6, 3)) - b = torch.ones((2)) + b = torch.ones(2) tensor_dataclass = TestTensorDataclass(a=a, b=b) assert tensor_dataclass.b.shape == (4, 6, 2) diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py index 0430ccc68..5a3c31300 100644 --- a/tests/viewer/test_viewer.py +++ b/tests/viewer/test_viewer.py @@ -71,7 +71,7 @@ def test_drawing(): v["random"].set_transform(get_translation_matrix([-0.5, -0.5, 0])) v = vis["lines"] - v.set_transform(get_translation_matrix(([-2, -3, 0]))) + v.set_transform(get_translation_matrix([-2, -3, 0])) vertices = np.random.random((3, 10)).astype(np.float32) v["line_segments"].set_object(g.LineSegments(g.PointsGeometry(vertices))) @@ -104,7 +104,7 @@ def test_drawing(): g.LineBasicMaterial(vertexColors=True), ) ) - v["triad"].set_transform(get_translation_matrix(([0, 5, 0]))) + v["triad"].set_transform(get_translation_matrix([0, 5, 0])) v["triad_function"].set_object(g.triad(0.5)) v["triad_function"].set_transform(get_translation_matrix([0, 6, 0]))