Small optimizations when working with images (#375)

* data optimization

* lint

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]:

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: