From f6108aaedefe8fd137b45fda338aae72b45f20a2 Mon Sep 17 00:00:00 2001 From: Jonatan Nevo Date: Wed, 8 Jan 2025 22:25:27 +0200 Subject: [PATCH] add 3d object rendering using open3d --- requirements.txt | 3 +- tagstudio/src/core/open3d_renderer.py | 141 +++++++++++++++++++++ tagstudio/src/core/utils/singleton.py | 13 ++ tagstudio/src/qt/widgets/thumb_renderer.py | 14 +- 4 files changed, 168 insertions(+), 3 deletions(-) create mode 100644 tagstudio/src/core/open3d_renderer.py create mode 100644 tagstudio/src/core/utils/singleton.py diff --git a/requirements.txt b/requirements.txt index 9eb290390..1232fd6f7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,4 +16,5 @@ SQLAlchemy==2.0.34 structlog==24.4.0 typing_extensions>=3.10.0.0,<=4.11.0 ujson>=5.8.0,<=5.9.0 -vtf2img==0.1.0 \ No newline at end of file +vtf2img==0.1.0 +open3d>=0.19.0, <1.0.0 \ No newline at end of file diff --git a/tagstudio/src/core/open3d_renderer.py b/tagstudio/src/core/open3d_renderer.py new file mode 100644 index 000000000..8544fcc95 --- /dev/null +++ b/tagstudio/src/core/open3d_renderer.py @@ -0,0 +1,141 @@ +import hashlib +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Optional +from queue import Queue, Empty + +import cv2 +import numpy as np +import open3d +from PIL import Image + +from core.utils.singleton import Singleton + +@dataclass +class QueueRequest: + filename: Path + size: tuple[int, int] + mesh: open3d.geometry.TriangleMesh + + +@dataclass +class QueueResponse: + filename: Path + size: tuple[int, int] + image: np.array + + +QUEUE_TIMEOUT = 0.1 +DEFAULT_SIZE = (160, 160) + +# Caching render results in memory, can be persistent if needed +class RenderCache: + def __init__(self): + self._cache: dict[str, Image.Image] = {} + self._lock = threading.Lock() + + @staticmethod + def _generate_key(filename: Path, size: tuple[int, int]) -> str: + with filename.open("rb") as f: + content_hash = hashlib.file_digest(f, "md5") + content_hash.update(str(size).encode()) + return content_hash.hexdigest() + + def get(self, filename: Path, size: tuple[int, int]) -> Optional[Image.Image]: + key = self._generate_key(filename, size) + with self._lock: + if key in self._cache: + return self._cache[key] + return None + + def put(self, filename: Path, size: tuple[int, int], image: Image.Image): + key = self._generate_key(filename, size) + with self._lock: + self._cache[key] = image + + def clear(self): + with self._lock: + self._cache.clear() + +# A thread safe class to handle multiple rendering calls to the Open3D library +class Open3DRenderer(metaclass=Singleton): + def __init__(self, render_cache: Optional[RenderCache] = None): + if render_cache is not None: + self._rendered_thumbnails_cache = render_cache + else: + self._rendered_thumbnails_cache = RenderCache() + + self._stop_event = threading.Event() + self._render_request_queue = Queue() + self._image_response_queue = Queue() + self._render_thread = threading.Thread(target=self._render_loop) + self._render_thread.start() + + + def render(self, filename: Path, size: tuple[int, int]) -> Image.Image: + image = self._rendered_thumbnails_cache.get(filename, size) + if image is None: + image = self._render(filename, size) + self._rendered_thumbnails_cache.put(filename, size, image) + return image + + def _render(self, filename: Path, size: tuple[int, int]) -> Image.Image: + mesh = open3d.io.read_triangle_mesh(filename) + mesh.compute_vertex_normals() + request = QueueRequest(filename, size, mesh) + 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 + + image_np = (response.image * 255).astype(np.uint8) + image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR) + return Image.fromarray(image_bgr) + + def _render_loop(self): + vis = open3d.visualization.Visualizer() + old_size = DEFAULT_SIZE + self._create_window(vis, old_size) + + while not self._stop_event.set(): + try: + request: QueueRequest = self._render_request_queue.get(timeout=QUEUE_TIMEOUT) + except Empty: + continue + + if request.size != old_size: + print(f"RESIZING from {old_size} to {request.size}") + vis.destroy_window() + self._create_window(vis, request.size) + old_size = request.size + + + vis.add_geometry(request.mesh) + + vis.poll_events() + vis.update_renderer() + + image = vis.capture_screen_float_buffer(do_render=True) + image_np = np.asarray(image) + + vis.remove_geometry(request.mesh) + + response = QueueResponse(request.filename, request.size, image_np) + self._image_response_queue.put(response) + + @staticmethod + def _create_window(visualiser: open3d.visualization.Visualizer, size: tuple[int, int]): + visualiser.create_window(visible=False, width=size[0], height=size[1]) + opt = visualiser.get_render_option() + opt.background_color = np.asarray([1, 1, 1]) + opt.light_on = True + opt.mesh_show_back_face = True diff --git a/tagstudio/src/core/utils/singleton.py b/tagstudio/src/core/utils/singleton.py new file mode 100644 index 000000000..76b758431 --- /dev/null +++ b/tagstudio/src/core/utils/singleton.py @@ -0,0 +1,13 @@ +# from gist: https://gist.github.com/JonatanNevo/c48efb9a13636252ddf48e3b864899f0 +from typing import TypeVar, Generic, Any + +T = TypeVar("T") + + +class Singleton(type, Generic[T]): + _instances: dict['Singleton[T]', T] = {} + + def __call__(cls, *args: Any, **kwargs: Any) -> T: + if cls not in cls._instances: + cls._instances[cls] = super().__call__(*args, **kwargs) + return cls._instances[cls] \ No newline at end of file diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py index b17f519ad..001d4f595 100644 --- a/tagstudio/src/qt/widgets/thumb_renderer.py +++ b/tagstudio/src/qt/widgets/thumb_renderer.py @@ -1,17 +1,19 @@ # Copyright (C) 2024 Travis Abendshien (CyanVoxel). # Licensed under the GPL-3.0 License. # Created for TagStudio: https://github.com/CyanVoxel/TagStudio - - +import functools +import hashlib import math import struct import zipfile from copy import deepcopy +from functools import lru_cache from io import BytesIO from pathlib import Path import cv2 import numpy as np +import open3d import pillow_jxl # noqa: F401 import rawpy import structlog @@ -44,6 +46,8 @@ from PySide6.QtGui import QGuiApplication, QImage, QPainter, QPixmap from PySide6.QtPdf import QPdfDocument, QPdfDocumentRenderOptions from PySide6.QtSvg import QSvgRenderer + +from core.open3d_renderer import Open3DRenderer from src.core.constants import FONT_SAMPLE_SIZES, FONT_SAMPLE_TEXT from src.core.exceptions import NoRendererError from src.core.media_types import MediaCategories, MediaType @@ -617,6 +621,9 @@ def _source_engine(self, filepath: Path) -> Image.Image: logger.error("Couldn't render thumbnail", filepath=filepath, error=type(e).__name__) return im + def _render_using_open3d(self, filepath: Path, size: tuple[int, int]) -> Image.Image: + return Open3DRenderer().render(filepath, size) + @classmethod def _open_doc_thumb(cls, filepath: Path) -> Image.Image: """Extract and render a thumbnail for an OpenDocument file. @@ -1143,6 +1150,9 @@ def render_unlinked() -> Image.Image: ext, MediaCategories.SOURCE_ENGINE_TYPES, mime_fallback=True ): image = self._source_engine(_filepath) + # Model Files ================================================== + elif MediaCategories.is_ext_in_category(ext, MediaCategories.MODEL_TYPES, mime_fallback=True): + image = self._render_using_open3d(_filepath, (adj_size, adj_size)) # No Rendered Thumbnail ======================================== if not image: raise NoRendererError