diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py index d30869114..13eb520be 100644 --- a/nerfactory/cameras/cameras.py +++ b/nerfactory/cameras/cameras.py @@ -18,11 +18,11 @@ Camera Models import base64 import math from enum import Enum, auto -from typing import Dict, Optional, Tuple, Union +from typing import Dict, Optional, Union import cv2 -import numpy as np import torch +import torchvision from torch.nn.functional import normalize from torchtyping import TensorType @@ -246,14 +246,14 @@ class Cameras: self, camera_idx: int, image: Optional[TensorType["height", "width", 2]] = None, - resize_shape: Optional[Tuple[int, int]] = None, + max_size: Optional[int] = None, ) -> Dict: """Convert a camera to a json dictionary. Args: camera_idx (int): Index of the camera to convert. image: An image in range [0, 1] that is encoded to a base64 string. Defaults to None. - resize_shape: Shape to resize the image to. Defaults to None. + max_size: Max size to resize the image to. Defaults to None. Returns: A JSON representation of the camera @@ -268,11 +268,14 @@ class Cameras: "camera_index": camera_idx, } if image is not None: - image_uint8 = (image * 255).detach().cpu().numpy().astype(np.uint8) - if resize_shape: - image_uint8 = cv2.resize(image_uint8, resize_shape) - data = cv2.imencode(".png", image_uint8)[1].tobytes() - json_["image"] = str("data:image/png;base64," + base64.b64encode(data).decode("ascii")) + image_uint8 = (image * 255).detach().type(torch.uint8) + if max_size is not None: + image_uint8 = image_uint8.permute(2, 0, 1) + image_uint8 = torchvision.transforms.functional.resize(image_uint8, max_size) # type: ignore + image_uint8 = image_uint8.permute(1, 2, 0) + image_uint8 = image_uint8.cpu().numpy() + data = cv2.imencode(".jpg", image_uint8)[1].tobytes() + json_["image"] = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii")) return json_ def get_intrinsics_matrices(self) -> TensorType["num_cameras", 3, 3]: diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 44bda7617..4794c8cf3 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -249,7 +249,7 @@ class VisualizerState: for idx in image_indices: image = dataset[idx]["image"] bgr = image[..., [2, 1, 0]] - camera_json = dataset.dataset_inputs.cameras.to_json(camera_idx=idx, image=bgr, resize_shape=(100, 100)) + camera_json = dataset.dataset_inputs.cameras.to_json(camera_idx=idx, image=bgr, max_size=100) self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json) # draw the scene bounds (i.e., the bounding box) @@ -390,9 +390,8 @@ class VisualizerState: self.output_type_changed = False self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type) self.vis["renderingState/colormap_options"].write(colormap_options) - selected_output = self._apply_colormap(outputs, stuff_colors).cpu().numpy() - image_output = selected_output * 255 - image = (image_output).astype("uint8") + selected_output = (self._apply_colormap(outputs, stuff_colors) * 255).type(torch.uint8) + image = selected_output.cpu().numpy() self.vis.set_image(image) def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int) -> None: