Skip to content

Commit

Permalink
Add RAW image support (Resolve #193)
Browse files Browse the repository at this point in the history 
- 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
  • Loading branch information
@CyanVoxel
CyanVoxel committed May 19, 2024
1 parent 699ecd3 commit 18becd6
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 33 deletions.
1 change: 1 addition & 0 deletions requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions tagstudio/src/core/constants.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
"j2k",
"jpg2",
]
RAW_IMAGE_TYPES: list[str] = ["raw", "dng", "rw2", "nef", "arw", "crw", "cr3"]
VIDEO_TYPES: list[str] = [
"mp4",
"webm",
Expand Down
34 changes: 13 additions & 21 deletions tagstudio/src/qt/widgets/preview_panel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
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
Expand All @@ -29,7 +30,7 @@

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
Expand Down Expand Up @@ -503,40 +504,28 @@ def update_widgets(self):
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"
)
else:
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 (
Expand All @@ -545,6 +534,9 @@ def update_widgets(self):
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})"
)
Expand Down
64 changes: 52 additions & 12 deletions tagstudio/src/qt/widgets/thumb_renderer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from pathlib import Path

import cv2
import rawpy
from PIL import (
Image,
UnidentifiedImageError,
Expand All @@ -22,7 +23,12 @@
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

Expand All @@ -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()

Expand All @@ -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),
)

Expand All @@ -98,9 +102,7 @@ def render(
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),
)

Expand Down Expand Up @@ -135,6 +137,25 @@ def render(
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)
Expand All @@ -160,6 +181,25 @@ def render(
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(
Expand Down

0 comments on commit 18becd6

Please sign in to comment.