feat: add pdf thumbnail support (port #378) (#543)

* feat: add pdf thumbnail support

Co-Authored-By: Heiholf <[email protected]>

* fix: remove redef

* tests: add test comparing pdf to png snapshot

Co-Authored-By: yed <[email protected]>

* fix: fix info in docstrings

* fix: remove sample png generation

* fix: change the pdf snapshot to use a black square

* chore: fix whitespace

---------

Co-authored-by: Heiholf <[email protected]>
Co-authored-by: yed <[email protected]>

tagstudio/src/qt/helpers/image_effects.py

@@ -0,0 +1,27 @@
+# Copyright (C) 2024 Travis Abendshien (CyanVoxel).
+# Licensed under the GPL-3.0 License.
+# Created for TagStudio: https://github.com/CyanVoxel/TagStudio
+
+import numpy as np
+from PIL import Image
+
+
+def replace_transparent_pixels(
+    img: Image.Image, color: tuple[int, int, int, int] = (255, 255, 255, 255)
+) -> Image.Image:
+    """Replace (copying/without mutating) all transparent pixels in an image with the color.
+
+    Args:
+        img (Image.Image):
+            The source image
+        color (tuple[int, int, int, int]):
+            The color (RGBA, 0 to 255) which transparent pixels should be set to.
+            Defaults to white (255, 255, 255, 255)
+
+    Returns:
+        Image.Image:
+            A copy of img with the pixels replaced.
+    """
+    pixel_array = np.asarray(img.convert("RGBA")).copy()
+    pixel_array[pixel_array[:, :, 3] == 0] = color
+    return Image.fromarray(pixel_array)

tagstudio/src/qt/widgets/thumb_renderer.py

@@ -28,8 +28,19 @@
 from PIL.Image import DecompressionBombError
 from pillow_heif import register_avif_opener, register_heif_opener
 from pydub import exceptions
-from PySide6.QtCore import QBuffer, QObject, QSize, Qt, Signal
+from PySide6.QtCore import (
+    QBuffer,
+    QFile,
+    QFileDevice,
+    QIODeviceBase,
+    QObject,
+    QSize,
+    QSizeF,
+    Qt,
+    Signal,
+)
 from PySide6.QtGui import QGuiApplication, QImage, QPainter, QPixmap
+from PySide6.QtPdf import QPdfDocument, QPdfDocumentRenderOptions
 from PySide6.QtSvg import QSvgRenderer
 from src.core.constants import FONT_SAMPLE_SIZES, FONT_SAMPLE_TEXT
 from src.core.media_types import MediaCategories, MediaType
@@ -39,6 +50,7 @@
 from src.qt.helpers.color_overlay import theme_fg_overlay
 from src.qt.helpers.file_tester import is_readable_video
 from src.qt.helpers.gradient import four_corner_gradient
+from src.qt.helpers.image_effects import replace_transparent_pixels
 from src.qt.helpers.text_wrapper import wrap_full_text
 from src.qt.helpers.vendored.pydub.audio_segment import (  # type: ignore
     _AudioSegment as AudioSegment,
@@ -812,6 +824,52 @@ def _model_stl_thumb(self, filepath: Path, size: int) -> Image.Image:
 
         return im
 
+    def _pdf_thumb(self, filepath: Path, size: int) -> Image.Image:
+        """Render a thumbnail for a PDF file.
+
+        filepath (Path): The path of the file.
+            size (int): The size of the icon.
+        """
+        im: Image.Image = None
+
+        file: QFile = QFile(filepath)
+        success: bool = file.open(
+            QIODeviceBase.OpenModeFlag.ReadOnly, QFileDevice.Permission.ReadUser
+        )
+        if not success:
+            logger.error("Couldn't render thumbnail", filepath=filepath)
+            return im
+        document: QPdfDocument = QPdfDocument()
+        document.load(file)
+        # Transform page_size in points to pixels with proper aspect ratio
+        page_size: QSizeF = document.pagePointSize(0)
+        ratio_hw: float = page_size.height() / page_size.width()
+        if ratio_hw >= 1:
+            page_size *= size / page_size.height()
+        else:
+            page_size *= size / page_size.width()
+        # Enlarge image for antialiasing
+        scale_factor = 2.5
+        page_size *= scale_factor
+        # Render image with no anti-aliasing for speed
+        render_options: QPdfDocumentRenderOptions = QPdfDocumentRenderOptions()
+        render_options.setRenderFlags(
+            QPdfDocumentRenderOptions.RenderFlag.TextAliased
+            | QPdfDocumentRenderOptions.RenderFlag.ImageAliased
+            | QPdfDocumentRenderOptions.RenderFlag.PathAliased
+        )
+        # Convert QImage to PIL Image
+        qimage: QImage = document.render(0, page_size.toSize(), render_options)
+        buffer: QBuffer = QBuffer()
+        buffer.open(QBuffer.OpenModeFlag.ReadWrite)
+        try:
+            qimage.save(buffer, "PNG")
+            im = Image.open(BytesIO(buffer.buffer().data()))
+        finally:
+            buffer.close()
+        # Replace transparent pixels with white (otherwise Background defaults to transparent)
+        return replace_transparent_pixels(im)
+
     def _text_thumb(self, filepath: Path) -> Image.Image:
         """Render a thumbnail for a plaintext file.
 
@@ -959,17 +1017,17 @@ def render(
                     else:
                         image = self._image_thumb(_filepath)
                 # Videos =======================================================
-                if MediaCategories.is_ext_in_category(
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.VIDEO_TYPES, mime_fallback=True
                 ):
                     image = self._video_thumb(_filepath)
                 # Plain Text ===================================================
-                if MediaCategories.is_ext_in_category(
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.PLAINTEXT_TYPES, mime_fallback=True
                 ):
                     image = self._text_thumb(_filepath)
                 # Fonts ========================================================
-                if MediaCategories.is_ext_in_category(
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.FONT_TYPES, mime_fallback=True
                 ):
                     if is_grid_thumb:
@@ -979,23 +1037,26 @@ def render(
                         # Large (Full Alphabet) Preview
                         image = self._font_long_thumb(_filepath, adj_size)
                 # Audio ========================================================
-                if MediaCategories.is_ext_in_category(
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.AUDIO_TYPES, mime_fallback=True
                 ):
                     image = self._audio_album_thumb(_filepath, ext)
                     if image is None:
                         image = self._audio_waveform_thumb(_filepath, ext, adj_size, pixel_ratio)
                         if image is not None:
                             image = self._apply_overlay_color(image, UiColor.GREEN)
-
-                # Blender ===========================================================
-                if MediaCategories.is_ext_in_category(
+                # Blender ======================================================
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.BLENDER_TYPES, mime_fallback=True
                 ):
                     image = self._blender(_filepath)
-
+                # PDF ==========================================================
+                elif MediaCategories.is_ext_in_category(
+                    ext, MediaCategories.PDF_TYPES, mime_fallback=True
+                ):
+                    image = self._pdf_thumb(_filepath, adj_size)
                 # VTF ==========================================================
-                if MediaCategories.is_ext_in_category(
+                elif MediaCategories.is_ext_in_category(
                     ext, MediaCategories.SOURCE_ENGINE_TYPES, mime_fallback=True
                 ):
                     image = self._source_engine(_filepath)

tagstudio/tests/qt/test_thumb_renderer.py

@@ -10,6 +10,17 @@
 from syrupy.extensions.image import PNGImageSnapshotExtension
 
 
+def test_pdf_preview(cwd, snapshot):
+    file_path: Path = cwd / "fixtures" / "sample.pdf"
+    renderer = ThumbRenderer()
+    img: Image.Image = renderer._pdf_thumb(file_path, 200)
+
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format="PNG")
+    img_bytes.seek(0)
+    assert img_bytes.read() == snapshot(extension_class=PNGImageSnapshotExtension)
+
+
 def test_svg_preview(cwd, snapshot):
     file_path: Path = cwd / "fixtures" / "sample.svg"
     renderer = ThumbRenderer()