From 18becd62a308e727b583caa387611f7fe453e2e4 Mon Sep 17 00:00:00 2001 From: Travis Abendshien Date: Sat, 18 May 2024 18:49:35 -0700 Subject: [PATCH] Add RAW image support (Resolve #193) - Add thumbnail and preview support for RAW images ["raw", "dng", "rw2", "nef", "arw", "crw", "cr3"] - Optimize the preview panel's dimension calculations (still need to move this elsewhere) - Refactored use of "Path" in thumb_renderer.py --- requirements.txt | 1 + tagstudio/src/core/constants.py | 1 + tagstudio/src/qt/widgets/preview_panel.py | 34 +++++------- tagstudio/src/qt/widgets/thumb_renderer.py | 64 ++++++++++++++++++---- 4 files changed, 67 insertions(+), 33 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0456b920..43c8a566 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ PySide6_Addons>=6.5.1.1,<=6.6.3.1 PySide6_Essentials>=6.5.1.1,<=6.6.3.1 typing_extensions>=3.10.0.0,<=4.11.0 ujson>=5.8.0,<=5.9.0 +rawpy==0.21.0 diff --git a/tagstudio/src/core/constants.py b/tagstudio/src/core/constants.py index 9e4eea76..8ae3cd1a 100644 --- a/tagstudio/src/core/constants.py +++ b/tagstudio/src/core/constants.py @@ -29,6 +29,7 @@ IMAGE_TYPES: list[str] = [ "j2k", "jpg2", ] +RAW_IMAGE_TYPES: list[str] = ["raw", "dng", "rw2", "nef", "arw", "crw", "cr3"] VIDEO_TYPES: list[str] = [ "mp4", "webm", diff --git a/tagstudio/src/qt/widgets/preview_panel.py b/tagstudio/src/qt/widgets/preview_panel.py index 0ae8d73f..049115af 100644 --- a/tagstudio/src/qt/widgets/preview_panel.py +++ b/tagstudio/src/qt/widgets/preview_panel.py @@ -9,6 +9,7 @@ import typing from datetime import datetime as dt import cv2 +import rawpy from PIL import Image, UnidentifiedImageError from PIL.Image import DecompressionBombError from PySide6.QtCore import Signal, Qt, QSize @@ -29,7 +30,7 @@ from humanfriendly import format_size from src.core.enums import SettingItems, Theme from src.core.library import Entry, ItemType, Library -from src.core.constants import VIDEO_TYPES, IMAGE_TYPES +from src.core.constants import VIDEO_TYPES, IMAGE_TYPES, RAW_IMAGE_TYPES from src.qt.helpers.file_opener import FileOpenerLabel, FileOpenerHelper, open_file from src.qt.modals.add_field import AddFieldModal from src.qt.widgets.thumb_renderer import ThumbRenderer @@ -503,30 +504,21 @@ class PreviewPanel(QWidget): image = None if extension in IMAGE_TYPES: image = Image.open(filepath) - if image.mode == "RGBA": - new_bg = Image.new("RGB", image.size, color="#1e1e1e") - new_bg.paste(image, mask=image.getchannel(3)) - image = new_bg - if image.mode != "RGB": - image = image.convert(mode="RGB") + elif extension in RAW_IMAGE_TYPES: + with rawpy.imread(filepath) as raw: + rgb = raw.postprocess() + image = Image.new( + "L", (rgb.shape[1], rgb.shape[0]), color="black" + ) elif extension in VIDEO_TYPES: video = cv2.VideoCapture(filepath) - video.set( - cv2.CAP_PROP_POS_FRAMES, - (video.get(cv2.CAP_PROP_FRAME_COUNT) // 2), - ) + video.set(cv2.CAP_PROP_POS_FRAMES, 0) success, frame = video.read() - if not success: - # Depending on the video format, compression, and frame - # count, seeking halfway does not work and the thumb - # must be pulled from the earliest available frame. - video.set(cv2.CAP_PROP_POS_FRAMES, 0) - success, frame = video.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image = Image.fromarray(frame) # Stats for specific file types are displayed here. - if extension in (IMAGE_TYPES + VIDEO_TYPES): + if extension in (IMAGE_TYPES + VIDEO_TYPES + RAW_IMAGE_TYPES): self.dimensions_label.setText( f"{extension.upper()} • {format_size(os.stat(filepath).st_size)}\n{image.width} x {image.height} px" ) @@ -534,9 +526,6 @@ class PreviewPanel(QWidget): self.dimensions_label.setText(f"{extension.upper()}") if not image: - self.dimensions_label.setText( - f"{extension.upper()} • {format_size(os.stat(filepath).st_size)}" - ) raise UnidentifiedImageError except ( @@ -545,6 +534,9 @@ class PreviewPanel(QWidget): cv2.error, DecompressionBombError, ) as e: + self.dimensions_label.setText( + f"{extension.upper()} • {format_size(os.stat(filepath).st_size)}" + ) logging.info( f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" ) diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py index 8bb14468..8748616b 100644 --- a/tagstudio/src/qt/widgets/thumb_renderer.py +++ b/tagstudio/src/qt/widgets/thumb_renderer.py @@ -9,6 +9,7 @@ import os from pathlib import Path import cv2 +import rawpy from PIL import ( Image, UnidentifiedImageError, @@ -22,7 +23,12 @@ from PIL.Image import DecompressionBombError from PySide6.QtCore import QObject, Signal, QSize from PySide6.QtGui import QPixmap from src.qt.helpers.gradient import four_corner_gradient_background -from src.core.constants import PLAINTEXT_TYPES, VIDEO_TYPES, IMAGE_TYPES +from src.core.constants import ( + PLAINTEXT_TYPES, + VIDEO_TYPES, + IMAGE_TYPES, + RAW_IMAGE_TYPES, +) ImageFile.LOAD_TRUNCATED_IMAGES = True @@ -42,29 +48,27 @@ class ThumbRenderer(QObject): # updatedSize = Signal(QSize) thumb_mask_512: Image.Image = Image.open( - Path(f"{Path(__file__).parents[3]}/resources/qt/images/thumb_mask_512.png") + Path(__file__).parents[3] / "resources/qt/images/thumb_mask_512.png" ) thumb_mask_512.load() thumb_mask_hl_512: Image.Image = Image.open( - Path(f"{Path(__file__).parents[3]}/resources/qt/images/thumb_mask_hl_512.png") + Path(__file__).parents[3] / "resources/qt/images/thumb_mask_hl_512.png" ) thumb_mask_hl_512.load() thumb_loading_512: Image.Image = Image.open( - Path(f"{Path(__file__).parents[3]}/resources/qt/images/thumb_loading_512.png") + Path(__file__).parents[3] / "resources/qt/images/thumb_loading_512.png" ) thumb_loading_512.load() thumb_broken_512: Image.Image = Image.open( - Path(f"{Path(__file__).parents[3]}/resources/qt/images/thumb_broken_512.png") + Path(__file__).parents[3] / "resources/qt/images/thumb_broken_512.png" ) thumb_broken_512.load() thumb_file_default_512: Image.Image = Image.open( - Path( - f"{Path(__file__).parents[3]}/resources/qt/images/thumb_file_default_512.png" - ) + Path(__file__).parents[3] / "resources/qt/images/thumb_file_default_512.png" ) thumb_file_default_512.load() @@ -75,7 +79,7 @@ class ThumbRenderer(QObject): # TODO: Make dynamic font sized given different pixel ratios font_pixel_ratio: float = 1 ext_font = ImageFont.truetype( - Path(f"{Path(__file__).parents[3]}/resources/qt/fonts/Oxanium-Bold.ttf"), + Path(__file__).parents[3] / "resources/qt/fonts/Oxanium-Bold.ttf", math.floor(12 * font_pixel_ratio), ) @@ -98,9 +102,7 @@ class ThumbRenderer(QObject): if ThumbRenderer.font_pixel_ratio != pixel_ratio: ThumbRenderer.font_pixel_ratio = pixel_ratio ThumbRenderer.ext_font = ImageFont.truetype( - Path( - f"{Path(__file__).parents[3]}/resources/qt/fonts/Oxanium-Bold.ttf" - ), + Path(__file__).parents[3] / "resources/qt/fonts/Oxanium-Bold.ttf", math.floor(12 * ThumbRenderer.font_pixel_ratio), ) @@ -135,6 +137,25 @@ class ThumbRenderer(QObject): f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {filepath} (because of {e})" ) + elif extension in RAW_IMAGE_TYPES: + try: + with rawpy.imread(filepath) as raw: + rgb = raw.postprocess() + image = Image.frombytes( + "RGB", + (rgb.shape[1], rgb.shape[0]), + rgb, + decoder_name="raw", + ) + except DecompressionBombError as e: + logging.info( + f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {filepath} (because of {e})" + ) + except rawpy._rawpy.LibRawIOError: + logging.info( + f"[ThumbRenderer]{ERROR} Couldn't Render thumbnail for raw image {filepath}" + ) + # Videos ======================================================= elif extension in VIDEO_TYPES: video = cv2.VideoCapture(filepath) @@ -160,6 +181,25 @@ class ThumbRenderer(QObject): draw = ImageDraw.Draw(bg) draw.text((16, 16), text, file=(255, 255, 255)) image = bg + # 3D =========================================================== + # elif extension == 'stl': + # # Create a new plot + # matplotlib.use('agg') + # figure = plt.figure() + # axes = figure.add_subplot(projection='3d') + + # # Load the STL files and add the vectors to the plot + # your_mesh = mesh.Mesh.from_file(filepath) + + # poly_collection = mplot3d.art3d.Poly3DCollection(your_mesh.vectors) + # poly_collection.set_color((0,0,1)) # play with color + # scale = your_mesh.points.flatten() + # axes.auto_scale_xyz(scale, scale, scale) + # axes.add_collection3d(poly_collection) + # # plt.show() + # img_buf = io.BytesIO() + # plt.savefig(img_buf, format='png') + # image = Image.open(img_buf) # No Rendered Thumbnail ======================================== else: image = ThumbRenderer.thumb_file_default_512.resize(