diff --git a/tagstudio/src/qt/helpers/file_tester.py b/tagstudio/src/qt/helpers/file_tester.py new file mode 100644 index 00000000..36a48c2b --- /dev/null +++ b/tagstudio/src/qt/helpers/file_tester.py @@ -0,0 +1,26 @@ +# Copyright (C) 2024 Travis Abendshien (CyanVoxel). +# Licensed under the GPL-3.0 License. +# Created for TagStudio: https://github.com/CyanVoxel/TagStudio + + +import ffmpeg +from pathlib import Path + + +def is_readable_video(filepath: Path | str): + """Test if a video is in a readable format. Examples of unreadable videos + include files with undetermined codecs and DRM-protected content. + + Args: + filepath (Path | str): + """ + probe = ffmpeg.probe(Path(filepath)) + for stream in probe["streams"]: + if stream.get("codec_tag_string") in [ + "[0][0][0][0]", + "drma", + "drms", + "drmi", + ]: + return False + return True diff --git a/tagstudio/src/qt/widgets/collage_icon.py b/tagstudio/src/qt/widgets/collage_icon.py index b9234d7d..a344ce0b 100644 --- a/tagstudio/src/qt/widgets/collage_icon.py +++ b/tagstudio/src/qt/widgets/collage_icon.py @@ -25,6 +25,7 @@ from PySide6.QtCore import ( from src.core.library import Library from src.core.constants import DOC_TYPES, VIDEO_TYPES, IMAGE_TYPES +from src.qt.helpers.file_tester import is_readable_video ERROR = f"[ERROR]" @@ -112,30 +113,31 @@ class CollageIconRenderer(QObject): except DecompressionBombError as e: logging.info(f"[ERROR] One of the images was too big ({e})") elif filepath.suffix.lower() in VIDEO_TYPES: - video = cv2.VideoCapture(str(filepath)) - video.set( - cv2.CAP_PROP_POS_FRAMES, - (video.get(cv2.CAP_PROP_FRAME_COUNT) // 2), - ) - 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) + if is_readable_video(filepath): + video = cv2.VideoCapture(str(filepath), cv2.CAP_FFMPEG) + video.set( + cv2.CAP_PROP_POS_FRAMES, + (video.get(cv2.CAP_PROP_FRAME_COUNT) // 2), + ) success, frame = video.read() - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - with Image.fromarray(frame, mode="RGB") as pic: - if keep_aspect: - pic.thumbnail(size) - else: - pic = pic.resize(size) - if data_tint_mode and color: - pic = ImageChops.hard_light( - pic, Image.new("RGB", size, color) - ) - # collage.paste(pic, (y*thumb_size, x*thumb_size)) - self.rendered.emit(pic) + 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) + with Image.fromarray(frame, mode="RGB") as pic: + if keep_aspect: + pic.thumbnail(size) + else: + pic = pic.resize(size) + if data_tint_mode and color: + pic = ImageChops.hard_light( + pic, Image.new("RGB", size, color) + ) + # collage.paste(pic, (y*thumb_size, x*thumb_size)) + self.rendered.emit(pic) except (UnidentifiedImageError, FileNotFoundError): logging.info( f"\n{ERROR} Couldn't read {entry.path}{os.sep}{entry.filename}" diff --git a/tagstudio/src/qt/widgets/preview_panel.py b/tagstudio/src/qt/widgets/preview_panel.py index 9781475f..5d73edad 100644 --- a/tagstudio/src/qt/widgets/preview_panel.py +++ b/tagstudio/src/qt/widgets/preview_panel.py @@ -47,6 +47,7 @@ from src.qt.widgets.text_box_edit import EditTextBox from src.qt.widgets.text_line_edit import EditTextLine from src.qt.helpers.qbutton_wrapper import QPushButtonWrapper from src.qt.widgets.video_player import VideoPlayer +from src.qt.helpers.file_tester import is_readable_video # Only import for type checking/autocompletion, will not be imported at runtime. @@ -558,25 +559,27 @@ class PreviewPanel(QWidget): ): pass elif filepath.suffix.lower() in VIDEO_TYPES: - video = cv2.VideoCapture(str(filepath)) - if video.get(cv2.CAP_PROP_FRAME_COUNT) <= 0: - raise cv2.error("File is invalid or has 0 frames") - video.set(cv2.CAP_PROP_POS_FRAMES, 0) - success, frame = video.read() - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - image = Image.fromarray(frame) - if success: - self.preview_img.hide() - self.preview_vid.play( - filepath, QSize(image.width, image.height) + if is_readable_video(filepath): + video = cv2.VideoCapture(str(filepath), cv2.CAP_FFMPEG) + video.set( + cv2.CAP_PROP_POS_FRAMES, + (video.get(cv2.CAP_PROP_FRAME_COUNT) // 2), ) - self.resizeEvent( - QResizeEvent( - QSize(image.width, image.height), - QSize(image.width, image.height), + success, frame = video.read() + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + image = Image.fromarray(frame) + if success: + self.preview_img.hide() + self.preview_vid.play( + filepath, QSize(image.width, image.height) ) - ) - self.preview_vid.show() + self.resizeEvent( + QResizeEvent( + QSize(image.width, image.height), + QSize(image.width, image.height), + ) + ) + self.preview_vid.show() # Stats for specific file types are displayed here. if image and filepath.suffix.lower() in ( diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py index 5dc5c7aa..c01a5a16 100644 --- a/tagstudio/src/qt/widgets/thumb_renderer.py +++ b/tagstudio/src/qt/widgets/thumb_renderer.py @@ -5,7 +5,6 @@ import logging import math -import sys import cv2 import rawpy import numpy @@ -41,6 +40,7 @@ from src.core.constants import ( ) from src.core.utils.encoding import detect_char_encoding from src.qt.helpers.blender_thumbnailer import blend_thumb +from src.qt.helpers.file_tester import is_readable_video ImageFile.LOAD_TRUNCATED_IMAGES = True @@ -172,13 +172,8 @@ class ThumbRenderer(QObject): # Videos ======================================================= elif _filepath.suffix.lower() in VIDEO_TYPES: - video = cv2.VideoCapture(str(_filepath), cv2.CAP_FFMPEG) - # Stupid check to try and tell if the codec can be read. - # TODO: Find a way to intercept the native FFMPEG errors. - h = int(video.get(cv2.CAP_PROP_FOURCC)) - codec = h.to_bytes(4, byteorder=sys.byteorder).decode() - logging.info(f"{codec} - {h} - {video.getBackendName()}") - if h != 22: + if is_readable_video(_filepath): + video = cv2.VideoCapture(str(_filepath), cv2.CAP_FFMPEG) video.set( cv2.CAP_PROP_POS_FRAMES, (video.get(cv2.CAP_PROP_FRAME_COUNT) // 2), @@ -198,6 +193,8 @@ class ThumbRenderer(QObject): success, frame = video.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image = Image.fromarray(frame) + else: + image = self.thumb_file_default_512 # Plain Text =================================================== elif _filepath.suffix.lower() in PLAINTEXT_TYPES: