From 4284957b238285bd5e332d54946558ba6deb2ccb Mon Sep 17 00:00:00 2001 From: Lunarient <92484433+Lunarient@users.noreply.github.com> Date: Wed, 12 Mar 2025 23:08:58 -0400 Subject: [PATCH 1/4] feat: add model thumbnail rendering --- pyproject.toml | 1 + src/tagstudio/qt/helpers/model_thumbnailer.py | 147 ++++++++++++++++++ src/tagstudio/qt/widgets/thumb_renderer.py | 15 ++ 3 files changed, 163 insertions(+) create mode 100644 src/tagstudio/qt/helpers/model_thumbnailer.py diff --git a/pyproject.toml b/pyproject.toml index dda76018e..4fe57a86c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ "typing_extensions>=3.10.0.0,<4.11.0", "ujson>=5.8.0,<5.9.0", "vtf2img==0.1.0", + "open3d==0.19.0" ] [project.optional-dependencies] diff --git a/src/tagstudio/qt/helpers/model_thumbnailer.py b/src/tagstudio/qt/helpers/model_thumbnailer.py new file mode 100644 index 000000000..0bedfa882 --- /dev/null +++ b/src/tagstudio/qt/helpers/model_thumbnailer.py @@ -0,0 +1,147 @@ +import threading +from dataclasses import dataclass +from pathlib import Path +from queue import Empty, Queue + +import numpy as np +import structlog +from open3d.io import read_triangle_model # type: ignore +from open3d.visualization.rendering import ( # type: ignore + MaterialRecord, + OffscreenRenderer, + TriangleMeshModel, +) +from PIL import Image + +logger = structlog.get_logger(__name__) + + +@dataclass +class QueueRequest: + filename: Path + size: tuple[int, int] + model: TriangleMeshModel + + +@dataclass +class QueueResponse: + filename: Path + size: tuple[int, int] + image: np.array + + +QUEUE_TIMEOUT = 0.1 +DEFAULT_SIZE = (256, 256) + + +# A thread safe class to handle multiple rendering calls to the Open3D library +class Open3DRenderer: + def __init__(self): + self._stop_event = threading.Event() + self._render_request_queue = Queue() + self._image_response_queue = Queue() + self.UP = [0, 1, 0] + self.FOV = 60 + self.DISTANCE_SCALE = 1.0 + self.BG_COLOR = (0.5, 0.5, 0.5, 1.0) + self.renderer = None + + # Primarily for .STL files + self.default_mat = MaterialRecord() + self.default_mat.base_color = [1.0, 0.5, 0.0, 1.0] + self.default_mat.shader = "defaultLit" + + self._render_thread = threading.Thread(target=self._render_loop, daemon=True) + self._render_thread.start() + + # ! I do not know why .mtl's are getting passed here so I just kick them out for now + def render(self, filename: Path, size: tuple[int, int]) -> Image.Image: + if filename.suffix == ".mtl": + return None + return self._render(filename, size) + + def _render(self, filename: Path, size: tuple[int, int]) -> Image.Image: + model = read_triangle_model(filename) + request = QueueRequest(filename, size, model) + self._render_request_queue.put(request) + + response: QueueResponse | None = None + while response is None: + # Fetch only the correct response + try: + response: QueueResponse = self._image_response_queue.get(timeout=QUEUE_TIMEOUT) + if response.filename != filename: + self._image_response_queue.put(response) + response = None + except Empty: + continue + + return Image.fromarray(response.image) + + def _update_camera(self, renderer: OffscreenRenderer, model: TriangleMeshModel): + combined_bounding_box = None + # Iterate through all meshes to compute the combined bounding box + for mesh_model in model.meshes: + mesh = mesh_model.mesh + bounding_box = mesh.get_axis_aligned_bounding_box() + + if combined_bounding_box is None: + combined_bounding_box = bounding_box + else: + combined_bounding_box = combined_bounding_box + bounding_box + + # Get the center of the combined bounding box + center = combined_bounding_box.get_center() + + # Calculate the diagonal size of the bounding box + diagonal = combined_bounding_box.get_extent() + distance = np.linalg.norm(diagonal) * self.DISTANCE_SCALE + eye = center + np.array([1, 1, 1]) * distance / np.linalg.norm([1, 1, 1]) + + # Vertical offset helps center object in render better + vertical_offset = 0.4 + eye[1] += vertical_offset + renderer.setup_camera(self.FOV, center, eye, self.UP) + + def _render_loop(self): + old_size = DEFAULT_SIZE + while not self._stop_event.set(): + try: + request: QueueRequest = self._render_request_queue.get(timeout=QUEUE_TIMEOUT) + except Empty: + continue + + if self.renderer is not None and request.size != old_size: + logger.info(f"Releasing renderer for resize from {old_size} to {request.size}") + del self.renderer + self.renderer = None + + if self.renderer is None: + logger.info(f"RESIZING from {old_size} to {request.size}") + self.renderer = OffscreenRenderer(request.size[0], request.size[1]) + old_size = request.size + # Signal that the renderer is ready for use + # self.renderer_ready_event.set() + + # # Wait until renderer is ready for the first use + # self.renderer_ready_event.wait() + + # Setup Scene + self.renderer.scene.clear_geometry() + self.renderer.scene.add_model("model", request.model) + + # If stl paint the model + if request.filename.suffix == ".stl": + self.renderer.scene.update_material(self.default_mat) + + self.renderer.scene.set_background(self.BG_COLOR) + + # Update the camera position + self._update_camera(self.renderer, request.model) + + # Render the image + image = self.renderer.render_to_image() + image_np = np.asarray(image) + + response = QueueResponse(request.filename, request.size, image_np) + self._image_response_queue.put(response) diff --git a/src/tagstudio/qt/widgets/thumb_renderer.py b/src/tagstudio/qt/widgets/thumb_renderer.py index ec9c7715f..fce7d7bf9 100644 --- a/src/tagstudio/qt/widgets/thumb_renderer.py +++ b/src/tagstudio/qt/widgets/thumb_renderer.py @@ -63,6 +63,7 @@ from tagstudio.qt.helpers.gradient import four_corner_gradient from tagstudio.qt.helpers.image_effects import replace_transparent_pixels from tagstudio.qt.helpers.text_wrapper import wrap_full_text +from tagstudio.qt.helpers.model_thumbnailer import Open3DRenderer from tagstudio.qt.helpers.vendored.pydub.audio_segment import ( _AudioSegment as AudioSegment, ) @@ -84,6 +85,7 @@ class ThumbRenderer(QObject): """A class for rendering image and file thumbnails.""" rm: ResourceManager = ResourceManager() + open3d_renderer = Open3DRenderer() cache: CacheManager = CacheManager() updated = Signal(float, QPixmap, QSize, Path, str) updated_ratio = Signal(float) @@ -613,6 +615,14 @@ def _blender(self, filepath: Path) -> Image.Image: else: logger.error("Couldn't render thumbnail", filepath=filepath, error=type(e).__name__) return im + + def _3d_model(self, filepath: Path, size: tuple[int, int]) -> Image.Image: + im: Image.Image = None + try: + im = self.open3d_renderer.render(filepath, size) + except Exception as e: + logger.error("Couldn't render 3d model", path=filepath.name, error=type(e).__name__) + return im def _source_engine(self, filepath: Path) -> Image.Image: """This is a function to convert the VTF (Valve Texture Format) files to thumbnails. @@ -1325,6 +1335,11 @@ def _render( ext, MediaCategories.SOURCE_ENGINE_TYPES, mime_fallback=True ): image = self._source_engine(_filepath) + # Model ========================================================== + elif MediaCategories.is_ext_in_category( + ext, MediaCategories.MODEL_TYPES, mime_fallback=True + ): + image = self._3d_model(_filepath, (adj_size,adj_size)) # No Rendered Thumbnail ======================================== if not image: raise NoRendererError From 1e5e48150a8cbb69fecc2560214b3ed86439fbee Mon Sep 17 00:00:00 2001 From: Lunarient <92484433+Lunarient@users.noreply.github.com> Date: Wed, 12 Mar 2025 23:38:33 -0400 Subject: [PATCH 2/4] fix (model-thumbnails): run ruff to fix code formatting issues --- src/tagstudio/qt/widgets/thumb_renderer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/tagstudio/qt/widgets/thumb_renderer.py b/src/tagstudio/qt/widgets/thumb_renderer.py index fce7d7bf9..3c4318919 100644 --- a/src/tagstudio/qt/widgets/thumb_renderer.py +++ b/src/tagstudio/qt/widgets/thumb_renderer.py @@ -62,8 +62,8 @@ from tagstudio.qt.helpers.file_tester import is_readable_video from tagstudio.qt.helpers.gradient import four_corner_gradient from tagstudio.qt.helpers.image_effects import replace_transparent_pixels -from tagstudio.qt.helpers.text_wrapper import wrap_full_text from tagstudio.qt.helpers.model_thumbnailer import Open3DRenderer +from tagstudio.qt.helpers.text_wrapper import wrap_full_text from tagstudio.qt.helpers.vendored.pydub.audio_segment import ( _AudioSegment as AudioSegment, ) @@ -615,7 +615,7 @@ def _blender(self, filepath: Path) -> Image.Image: else: logger.error("Couldn't render thumbnail", filepath=filepath, error=type(e).__name__) return im - + def _3d_model(self, filepath: Path, size: tuple[int, int]) -> Image.Image: im: Image.Image = None try: @@ -1339,7 +1339,7 @@ def _render( elif MediaCategories.is_ext_in_category( ext, MediaCategories.MODEL_TYPES, mime_fallback=True ): - image = self._3d_model(_filepath, (adj_size,adj_size)) + image = self._3d_model(_filepath, (adj_size, adj_size)) # No Rendered Thumbnail ======================================== if not image: raise NoRendererError From 88c4ac4de789a2b677ee993bb2e4d75a7fb6f969 Mon Sep 17 00:00:00 2001 From: Lunarient <92484433+Lunarient@users.noreply.github.com> Date: Thu, 13 Mar 2025 00:00:12 -0400 Subject: [PATCH 3/4] fix (model-thumbnails): mypy errors --- src/tagstudio/qt/helpers/model_thumbnailer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/tagstudio/qt/helpers/model_thumbnailer.py b/src/tagstudio/qt/helpers/model_thumbnailer.py index 0bedfa882..26c50cc58 100644 --- a/src/tagstudio/qt/helpers/model_thumbnailer.py +++ b/src/tagstudio/qt/helpers/model_thumbnailer.py @@ -4,9 +4,10 @@ from queue import Empty, Queue import numpy as np +import numpy.typing as npt import structlog -from open3d.io import read_triangle_model # type: ignore -from open3d.visualization.rendering import ( # type: ignore +from open3d.io import read_triangle_model +from open3d.visualization.rendering import ( MaterialRecord, OffscreenRenderer, TriangleMeshModel, @@ -27,7 +28,7 @@ class QueueRequest: class QueueResponse: filename: Path size: tuple[int, int] - image: np.array + image: npt.NDArray QUEUE_TIMEOUT = 0.1 @@ -69,7 +70,7 @@ def _render(self, filename: Path, size: tuple[int, int]) -> Image.Image: while response is None: # Fetch only the correct response try: - response: QueueResponse = self._image_response_queue.get(timeout=QUEUE_TIMEOUT) + response = self._image_response_queue.get(timeout=QUEUE_TIMEOUT) if response.filename != filename: self._image_response_queue.put(response) response = None From 2a4f30aeba2992c2d1c8b6a8b7a58d7375404b1f Mon Sep 17 00:00:00 2001 From: Lunarient <92484433+Lunarient@users.noreply.github.com> Date: Thu, 13 Mar 2025 00:08:55 -0400 Subject: [PATCH 4/4] refactor (model-thumbnails): remove unused lines --- src/tagstudio/qt/helpers/model_thumbnailer.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/tagstudio/qt/helpers/model_thumbnailer.py b/src/tagstudio/qt/helpers/model_thumbnailer.py index 26c50cc58..0bdeaa2d5 100644 --- a/src/tagstudio/qt/helpers/model_thumbnailer.py +++ b/src/tagstudio/qt/helpers/model_thumbnailer.py @@ -121,11 +121,6 @@ def _render_loop(self): logger.info(f"RESIZING from {old_size} to {request.size}") self.renderer = OffscreenRenderer(request.size[0], request.size[1]) old_size = request.size - # Signal that the renderer is ready for use - # self.renderer_ready_event.set() - - # # Wait until renderer is ready for the first use - # self.renderer_ready_event.wait() # Setup Scene self.renderer.scene.clear_geometry()