add quality setting; allow png for lossless compression (#1639)

* add quality setting; allow png for lossless compression

* black reformatting

* JPEG default, 90 quality

nerfstudio/configs/base_config.py

@@ -23,11 +23,14 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, List, Optional, Tuple, Type
 
+from typing_extensions import Literal
+
 # model instances
 from nerfstudio.utils import writer
 
 warnings.filterwarnings("ignore", module="torchvision")
 
+
 # Pretty printing class
 class PrintableConfig:  # pylint: disable=too-few-public-methods
     """Printable Config defining str function"""
@@ -148,3 +151,9 @@ class ViewerConfig(PrintableConfig):
     actually used in training/evaluation. If -1, display all."""
     quit_on_train_completion: bool = False
     """Whether to kill the training job when it has completed. Note this will stop rendering in the viewer."""
+    image_format: Literal["jpeg", "png"] = "jpeg"
+    """Image format viewer should use; jpeg is lossy compression, while png is lossless."""
+    jpeg_quality: int = 90
+    """Quality tradeoff to use for jpeg compression."""
+    png_compression: int = 1
+    """Size/speed tradeoff to use for png compression."""

nerfstudio/viewer/server/viewer_utils.py

@@ -670,8 +670,17 @@ class ViewerState:
 
         image = selected_output[..., [2, 1, 0]].cpu().numpy()
 
-        data = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 75])[1].tobytes()
-        data = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii"))
+        data = cv2.imencode(
+            f".{self.config.image_format}",
+            image,
+            [
+                cv2.IMWRITE_JPEG_QUALITY,
+                self.config.jpeg_quality,
+                cv2.IMWRITE_PNG_COMPRESSION,
+                self.config.png_compression,
+            ],
+        )[1].tobytes()
+        data = str(f"data:image/{self.config.image_format};base64," + base64.b64encode(data).decode("ascii"))
         self.vis["render_img"].write(data)
 
     def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int, image_width: int) -> None: