From 6db1d314cc646598daf0093b3a0513bbb43b718f Mon Sep 17 00:00:00 2001 From: machenmusik Date: Tue, 21 Mar 2023 19:04:12 -0400 Subject: [PATCH] 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 | 9 +++++++++ nerfstudio/viewer/server/viewer_utils.py | 13 +++++++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index fda6cc10f..3d7872d74 100644 --- a/nerfstudio/configs/base_config.py +++ b/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.""" diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index b781bea08..a6ff0a412 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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: