From 2f937b8f575471e3de586abc1e4a7ddfd64aeba5 Mon Sep 17 00:00:00 2001 From: LightArrowsEXE Date: Thu, 26 Sep 2024 14:53:45 +0200 Subject: [PATCH] npy_to_clip: Works now --- lvsfunc/exceptions.py | 17 +++++- lvsfunc/models/base.py | 3 +- lvsfunc/nn/func.py | 117 +++++++++++++++++++++++------------------ lvsfunc/nn/util.py | 8 +-- 4 files changed, 90 insertions(+), 55 deletions(-) diff --git a/lvsfunc/exceptions.py b/lvsfunc/exceptions.py index 9ed0c05..4f0434d 100644 --- a/lvsfunc/exceptions.py +++ b/lvsfunc/exceptions.py @@ -1,9 +1,14 @@ from __future__ import annotations -from vstools import CustomTypeError, FuncExceptT +from typing import Any + +from stgpytools import SupportsString +from vstools import CustomTypeError, CustomValueError, FuncExceptT __all__ = [ 'ClipsAndNamedClipsError', + + 'NumpyArrayLoadError', ] @@ -14,3 +19,13 @@ def __init__( self, func: FuncExceptT, message: str = 'Positional clips and named keyword clips cannot both be given!' ) -> None: super().__init__(message, func) + + +class NumpyArrayLoadError(CustomValueError): + """Raised when there's an issue with loading a numpy array.""" + + def __init__( + self, message: SupportsString | None = None, func: FuncExceptT | None = None, reason: Any = None, **kwargs: Any + ) -> None: + + super().__init__(message, func, reason, **kwargs) diff --git a/lvsfunc/models/base.py b/lvsfunc/models/base.py index 9beec54..64dfee6 100644 --- a/lvsfunc/models/base.py +++ b/lvsfunc/models/base.py @@ -8,7 +8,8 @@ from vsscale import autoselect_backend from vstools import (ColorRange, CustomValueError, DependencyNotFoundError, FileWasNotFoundError, FunctionUtil, Matrix, SPath, depth, - inject_self, iterate, join, normalize_planes, split, vs, get_peak_value) + get_peak_value, inject_self, iterate, join, + normalize_planes, split, vs) __all__: list[str] = [] diff --git a/lvsfunc/nn/func.py b/lvsfunc/nn/func.py index 1c8a37e..cbf0d28 100644 --- a/lvsfunc/nn/func.py +++ b/lvsfunc/nn/func.py @@ -2,10 +2,11 @@ import numpy as np from vsexprtools import norm_expr -from vstools import (CustomValueError, FuncExceptT, FunctionUtil, SPath, - SPathLike, clip_async_render, core, fallback, get_depth, - initialize_clip, vs) +from vstools import (CustomValueError, FileWasNotFoundError, FuncExceptT, + FunctionUtil, SPath, SPathLike, clip_async_render, core, + fallback, get_depth, initialize_clip, vs) +from ..exceptions import NumpyArrayLoadError from .util import get_format_from_npy __all__: list[str] = [ @@ -183,100 +184,116 @@ def npy_to_clip( :param func_except: Function returned for custom error handling. This should only be set by VS package developers. - :return: The numpy array as a clip. + :return: VapourSynth clip created from the numpy files. """ func = fallback(func_except, npy_to_clip) + paths = [SPath(file_paths)] if not isinstance(file_paths, list) else [SPath(x) for x in file_paths] is_npz = False - if not isinstance(file_paths, list): - file_paths_list = [SPath(file_paths)] - else: - file_paths_list = [SPath(fp) for fp in file_paths] - - if file_paths_list and file_paths_list[0].is_dir(): - file_paths_list = list(file_paths_list[0].glob("*.np[yz]")) + if not paths: + raise FileWasNotFoundError("No files provided!", func) - if not file_paths_list: - raise CustomValueError("No .npy or .npz files found in the directory", func) - elif not file_paths_list: - raise CustomValueError("No files provided", func) - - is_npz = file_paths_list[0].suffix == '.npz' + if paths[0].is_dir(): + if npy_files := list(paths[0].glob("*.npy")): + paths = npy_files + elif npz_files := list(paths[0].glob("*.npz")): + paths = npz_files + is_npz = True + else: + raise FileWasNotFoundError("No .npy or .npz files found in the given directory!", func) + elif paths[0].suffix == '.npz': + is_npz = True if not is_npz: - file_paths_list = sorted( - file_paths_list, key=lambda x: int(x.stem) if x.stem.isdigit() else float('inf') - ) + try: + paths = sorted(paths, key=lambda x: int(x.stem) if x.stem.isdigit() else float('inf')) + except ValueError as e: + if "invalid literal for int() with base 10" in str(e): + raise CustomValueError("Error sorting paths: File names must be valid integers!", func) + else: + raise + except Exception as e: + raise CustomValueError(f"Error sorting paths! {str(e)}", func) if is_npz: - with np.load(file_paths_list[0], allow_pickle=True) as npz_data: - first_frame = next(iter(npz_data.values())).item() + npz_data = np.load(paths[0], allow_pickle=True) + first_key = list(npz_data.keys())[0] + first_frame = npz_data[first_key] + + if isinstance(first_frame, np.ndarray) and first_frame.shape == (): + first_frame = first_frame.item() else: - first_frame = np.load(file_paths_list[0], allow_pickle=True).item() + first_frame = np.load(paths[0], allow_pickle=True).item() if isinstance(first_frame, dict): - plane_0 = first_frame.get('plane_0') + plane_0_missing = 'plane_0' not in first_frame + plane_0_not_array = not isinstance(first_frame['plane_0'], np.ndarray) + plane_0_low_dim = first_frame['plane_0'].ndim < 2 - if not isinstance(plane_0, np.ndarray) or plane_0.ndim < 2: - raise CustomValueError( - "Invalid frame data structure. 'plane_0' is missing, not a numpy array, " - "or has less than 2 dimensions!", func - ) + if plane_0_missing or plane_0_not_array or plane_0_low_dim: + error_message = "Invalid frame data structure." + + if plane_0_missing: + error_message += " 'plane_0' is missing from the frame dictionary." + elif plane_0_not_array: + error_message += " 'plane_0' is not a numpy array." + elif plane_0_low_dim: + error_message += " 'plane_0' has less than 2 dimensions." - height, width = plane_0.shape + raise NumpyArrayLoadError(error_message, func) + + height, width = first_frame['plane_0'].shape elif isinstance(first_frame, np.ndarray): if first_frame.ndim < 2: - raise CustomValueError( - f"Invalid frame shape: {first_frame.shape}. Expected at least 2 dimensions!", func + raise NumpyArrayLoadError( + f"Invalid frame shape: {first_frame.shape}. Expected at least 2 dimensions.", func ) height, width = first_frame.shape[:2] else: - raise CustomValueError(f"Unsupported frame data type: {type(first_frame)}", func) + raise NumpyArrayLoadError(f"Unsupported frame data type: {type(first_frame)}", func) fmt = get_format_from_npy(first_frame) + clip_length = len(npz_data.keys()) if is_npz else len(paths) blank_clip = core.std.BlankClip( - None, width, height, fmt, - length=len(npz_data.keys()) if is_npz else len(file_paths_list), - keep=True + None, width, height, fmt, length=clip_length, keep=True ) def _read_frame(n: int, f: vs.VideoFrame) -> vs.VideoFrame: if is_npz: - loaded_frame = next(iter(npz_data.values())) + loaded_frame = npz_data[list(npz_data.keys())[n]] + if isinstance(loaded_frame, np.ndarray) and loaded_frame.shape == (): + loaded_frame = loaded_frame.item() else: - loaded_frame = np.load(file_paths_list[n], allow_pickle=True) - - loaded_frame = loaded_frame.item() + loaded_frame = np.load(paths[n], allow_pickle=True).item() fout = f.copy() if isinstance(loaded_frame, np.ndarray): if loaded_frame.ndim < 2: - raise CustomValueError( - f"Invalid frame shape at index {n}: {loaded_frame.shape}. " - "Expected at least 2 dimensions!", func + raise NumpyArrayLoadError( + f"Invalid frame shape at index {n}: {loaded_frame.shape}. Expected at least 2 dimensions.", func ) + np.copyto(np.asarray(fout[0]), loaded_frame) elif isinstance(loaded_frame, dict): for plane in range(f.format.num_planes): plane_key = f'plane_{plane}' - plane_data = loaded_frame.get(plane_key) - if plane_data is None or plane_data.ndim < 2: - raise CustomValueError( + if plane_key not in loaded_frame or loaded_frame[plane_key].ndim < 2: + raise NumpyArrayLoadError( f"Invalid frame data structure at index {n}. '{plane_key}' " - "is missing or has less than 2 dimensions!", func=func + "is missing or has less than 2 dimensions.", func ) - np.copyto(np.asarray(fout[plane]), plane_data) + np.copyto(np.asarray(fout[plane]), loaded_frame[plane_key]) else: - raise CustomValueError(f"Unsupported frame data type at index {n}: {type(loaded_frame)}", func=func) + raise CustomValueError(f"Unsupported frame data type at index {n}: {type(loaded_frame)}", func) return fout proc_clip = blank_clip.std.ModifyFrame(blank_clip, _read_frame) - return initialize_clip(proc_clip, ref or get_depth(proc_clip), func=func) + return initialize_clip(proc_clip, ref or get_depth(proc_clip), func=func) # type:ignore[arg-type] diff --git a/lvsfunc/nn/util.py b/lvsfunc/nn/util.py index c8c7da0..fce39c9 100644 --- a/lvsfunc/nn/util.py +++ b/lvsfunc/nn/util.py @@ -1,6 +1,8 @@ import numpy as np -from vstools import (CustomTypeError, FuncExceptT, InvalidVideoFormatError, - core, depth, fallback, get_video_format, vs) +from vstools import (FuncExceptT, InvalidVideoFormatError, core, depth, + fallback, get_video_format, vs) + +from ..exceptions import NumpyArrayLoadError __all__: list[str] = [ 'get_format_from_npy', @@ -35,7 +37,7 @@ def get_format_from_npy(frame_data: np.ndarray, func_except: FuncExceptT | None y_data = frame_data num_planes = y_data.ndim if y_data.ndim <= 2 else y_data.shape[2] else: - raise CustomTypeError(f"Unsupported data type: {type(frame_data)}", func) + raise NumpyArrayLoadError(f"Unsupported data type: {type(frame_data)}", func) bit_depth = 32 if y_data.dtype == np.float32 else y_data.itemsize * 8