daio 0.0.1__py3-none-any.whl

daio/h5.py

@@ -0,0 +1,172 @@
+import json 
+#TODO: test replacing json with orjson
+import warnings
+from os.path import expanduser
+
+import numpy as np
+import h5py
+#import hdf5plugin
+#TODO: test defaulting to hdf5plugin.Zstd compression for json-serialized strings
+
+
+def save_to_h5(filename, data, serialize=True, compression=None, json_compression='gzip', verbosity=1, file_mode='w', convert_numpy_to_native=False):
+    ''' Save a nested dictionary data structure to an HDF5 file. 
+
+    Args:
+        filename (string): file name of the HDF5 file
+        data (dict): Nested dictionary whose contents may be dict, ndarray, str, bytes, DataFrame and JSON-serializable objects
+        serialize (boolean): enable JSON serialization
+        compression (string): h5py compression type (e.g. 'gzip', 'lzf' or None)
+        json_compression (string): h5py compression type for serialized JSON (default: 'gzip')
+        file_mode (string): h5py.File access mode. 'w' (default) for create/detete and 'a' for create/append
+
+    based on https://github.com/danionella/lib2p/blob/master/lib2putils.py
+    '''
+
+    def recursively_save_contents_to_group(h5file, path, data_item):
+        assert isinstance(data_item, (dict))
+        for key, item in data_item.items():
+            if verbosity > 1:
+                print('saving entry: {} -- {}'.format(path + key, type(item)))
+            if isinstance(item, (np.ndarray, np.int64, np.float64, str, bytes, int, float)):
+                comp = None if np.isscalar(item) else compression
+                try:
+                    h5file[path].create_dataset(key, data=item, compression=comp)
+                except TypeError:
+                    warnings.warn(f'\n\nkey: {key} -- Saving data with compression failed. Saving without compression.\n')
+                    h5file[path].create_dataset(key, data=item, compression=None)
+            elif type(item).__name__ == 'DataFrame':
+                json_bytes = np.frombuffer(item.to_json().encode('utf-8'), dtype='byte')
+                h5file[path].create_dataset(key, data=json_bytes, compression=json_compression)
+                h5file[path + key].attrs['pandas_json_type'] = f'This {type(item)} was JSON serialized and UTF-8 encoded.'
+            elif isinstance(item, dict):
+                h5file[path].create_group(key)
+                recursively_save_contents_to_group(h5file, path + key + '/', item)
+            elif serialize:
+                if verbosity > 0:
+                    print(f'serializing {type(item)} at {path+key}', flush=True)
+                #TODO: test replacing json with orjson
+                json_bytes = json.dumps(item).encode('utf-8')
+                h5file[path].create_dataset(key, data=np.frombuffer(json_bytes, dtype='byte'), compression=json_compression)
+                h5file[path + key].attrs['json_type'] = f'This {type(item)} was JSON serialized and UTF-8 encoded.'
+            else:
+                raise ValueError(f'Cannot save {type(item)} to {path+key}. Consider enabling serialisation.')
+
+    if convert_numpy_to_native:
+        data = convert_numpy_to_native(data)
+
+    filename = expanduser(filename)
+    with h5py.File(filename, file_mode) as h5file:
+        recursively_save_contents_to_group(h5file, '/', data)
+
+
+def load_from_h5(filename):
+    ''' Load an HDF5 file to a dictionary
+    
+    Args:
+        filename (string): file name of the HDF5 file
+        
+    Returns:
+        dict: file contents
+    '''
+
+    def recursively_load_contents_from_group(h5file, path):
+        ans = dict()
+        for key, item in h5file[path].items():
+            if 'pandas_type' in item.attrs.keys():
+                import pandas as pd
+                ans[key] = pd.read_hdf(filename, path + key)
+            elif 'pandas_json_type' in item.attrs.keys():
+                import pandas as pd
+                json_str = item[()].tobytes().decode('utf-8')
+                ans[key] = pd.read_json(json_str)
+            elif 'json_type' in item.attrs.keys():
+                #TODO: test replacing json with orjson
+                ans[key] = json.loads(item[()].tobytes())
+            elif isinstance(item, h5py._hl.dataset.Dataset):
+                if h5py.check_string_dtype(item.dtype) is not None:
+                    item = item.asstr()
+                ans[key] = item[()]
+            elif isinstance(item, h5py._hl.group.Group):
+                ans[key] = recursively_load_contents_from_group(h5file, path + key + '/')
+            else:
+                raise ValueError(f"I don't know what to do about {path+key}.")
+        return ans
+
+    filename = expanduser(filename)
+    with h5py.File(filename, 'r') as h5file:
+        return recursively_load_contents_from_group(h5file, '/')
+    
+
+class lazyh5:
+    """ A lazy-loading interface for HDF5 files. 
+    
+    This class provides an easy way to access HDF5 file content without fully
+    loading it into memory. It supports dynamic access to datasets and subgroups.
+
+    Args:
+            filepath (str): Path to the HDF5 file.
+            h5path (str, optional): HDF5 group path. Defaults to '/'.
+            overwrite (bool, optional): Whether to overwrite existing items. Defaults to False.
+    """
+
+    def __init__(self, filepath, h5path='/', overwrite=False):
+        self._filepath = filepath
+        self._h5path = h5path
+        self._overwrite = overwrite
+
+    def keys(self):
+        """Lists the keys of the current HDF5 group.
+
+        Returns:
+            list: List of keys in the current HDF5 group.
+        """
+        with h5py.File(self._filepath, 'r') as f:
+            return list(f[self._h5path].keys())
+
+    def __getitem__(self, key):
+        """Gets an item by key."""
+        with h5py.File(self._filepath, 'r') as f:
+            item = f[self._h5path][key]
+            if isinstance(item, h5py.Group):
+                return lazyh5(self._filepath, h5path=f"{self._h5path}/{key}")
+            elif isinstance(item, h5py.Dataset):
+                return item[()]
+            else:
+                raise KeyError(f"Unknown item type for key: {key}")
+
+    def __getattr__(self, key):
+        """Provides dynamic attribute-style access."""
+        try:
+            return self[key]
+        except KeyError:
+            raise AttributeError(key)
+
+    def __setattr__(self, key, value):
+        """Sets an attribute or creates a new dataset."""
+        if key.startswith('_'):
+            super().__setattr__(key, value)
+            return
+
+        with h5py.File(self._filepath, 'a') as f:
+            full_path = f"{self._h5path}/{key}".lstrip('/')
+            if (full_path in f) and (not self._overwrite):
+                raise AttributeError(f"Dataset or group '{key}' already exists.")
+            else:
+                f[self._h5path].create_dataset(key, data=value)
+
+    def __len__(self):
+        """Gets the number of items in the current HDF5 group."""
+        return len(self.keys())
+
+    def __repr__(self):
+        """Provides a string representation of the object."""
+        return f"<lazyh5 for file '{self._filepath}', HDF5 path '{self._h5path}' with {len(self)} items>"
+
+    def _ipython_key_completions_(self):
+        """Enables key completions in IPython."""
+        return self.keys()
+
+    def __dir__(self):
+        """Lists all accessible attributes and keys."""
+        return self.keys() + dir(super())

daio/video.py

@@ -0,0 +1,265 @@
+import warnings
+
+import numpy as np
+import av
+
+warnings.simplefilter("once")
+
+class VideoReader:
+    ''' A class to read video files quickly, using PyAV. It allows for fast seeking and reading of frames. 
+    
+    Args:
+        filename (str): path to the video file
+        format (str): format to read the video in. Default is 'gray'. Other options include 'rgb24', 'yuv420p', etc.
+        threading (bool): enable threading in the decoder. Default is True.
+        thread_count (int): number of threads to use for decoding. Default is 0 (auto).
+        tolerance (float): tolerance for frame time stamp deviations (fraction of frame interval). Default is 0.1.
+    '''
+    def __init__(self, filename, format='gray', threading=True, thread_count=0, tolerance=0.1):
+        self.container = av.open(filename)
+        self.stream = self.container.streams.video[0]
+        self.stream.codec_context.thread_count = thread_count
+        self.stream.codec_context.thread_type = 'AUTO' if threading else 'SLICE' # FRAME/AUTO/SLICE
+        self.framegenerator = self.container.decode(video=0)
+        self.read_format = format
+        self._pts_per_frame = 1 / (self.stream.guessed_rate * self.stream.time_base)
+        self._frame_to_pts = lambda n: round(n * self._pts_per_frame) + self.stream.start_time
+        self.wiggle_pts = tolerance * self._pts_per_frame
+        self.rewind()
+        if self.container.format.variable_fps:
+            warn_transcode(f'Variable frame rate video detected. Seeking will likely be unrealiable. I will warn again if I detect seek gitches')
+        if self.stream.average_rate != self.stream.guessed_rate:
+            warn_transcode(f'Average frame rate ({self.stream.average_rate}) is different from nominal frame rate ({self.stream.guessed_rate}). Seeking may be unrealiable. I will warn again if I detect seek gitches')
+
+    def read(self):
+        ''' Read the next frame in the specified format. '''
+        frame_obj = next(self.framegenerator)
+        self.last_pts = frame_obj.pts
+        im = frame_obj.to_ndarray(format=self.read_format)
+        return im
+
+    def rewind(self):
+        ''' Rewind the video to the beginning. '''
+        self.container.seek(0)
+        self.framegenerator = self.container.decode(video=0)
+        self.last_pts = None
+
+    def read_frame(self, frame_idx):
+        ''' Read the specified frame index. 
+        
+        Args:
+            frame_idx (int): index of the frame to read
+        '''
+        if frame_idx == 0:
+            self.rewind()
+            return self.read()
+        if self.last_pts is not None and (np.abs(self.last_pts - self._frame_to_pts(frame_idx-1)) <= self.wiggle_pts):
+            return self.read()
+        target_pts = self._frame_to_pts(frame_idx)
+        self.container.seek(target_pts-self.stream.start_time, backward=True, stream=self.container.streams.video[0])
+        self.framegenerator = self.container.decode(video=0)
+        frame_obj = next(self.framegenerator)
+        if frame_obj.pts > target_pts: #detecting overshoot (may happen due to variable frame rate)
+            n_back = 100
+            warn_transcode(f'Seek overshoot ({frame_obj.pts} > {target_pts}). Backtracking by {n_back} frames...')
+            self.container.seek(self._frame_to_pts(frame_idx-n_back)-self.stream.start_time, backward=True, stream=self.container.streams.video[0])
+            self.framegenerator = self.container.decode(video=0)
+            frame_obj = next(self.framegenerator)
+        while frame_obj.pts < (target_pts - self.wiggle_pts): 
+            frame_obj = next(self.framegenerator)
+        # frame_obj.pts should now be equal to target_pts
+        if np.abs(frame_obj.pts - target_pts) > self.wiggle_pts:
+            warn_transcode(f'Seek problem with frame {frame_idx}! pts: {frame_obj.pts}; target: {target_pts}; dts: {frame_obj.dts}; pict_type: {str(frame_obj.pict_type)}')
+        frame = frame_obj.to_ndarray(format=self.read_format)
+        self.last_pts = frame_obj.pts
+        return frame
+    
+    def close(self):
+        self.container.close()
+        
+    def __del__(self):
+        self.close()
+        
+    def __getitem__(self, index):
+        if isinstance(index, (int, np.integer)):  # single frame
+            return self.read_frame(index)
+        elif isinstance(index, slice):
+            frames = [self.__getitem__(i) for i in np.r_[index]]
+            return np.array(frames)
+        elif isinstance(index, tuple):
+            frames = [self.__getitem__(i)[index[1:]] for i in np.r_[index[0]]]
+            if len(frames) == 1:
+                return frames[0]
+            else:
+                return np.array(frames)
+        else:
+            raise NotImplementedError(f"slicing of {type(index)} : {index} not implemented yet")
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False
+    
+    def __iter__(self):
+        self.rewind()
+        return self
+    
+    def __next__(self):
+        try:
+            return self.read()
+        except:
+            self.rewind()
+            raise StopIteration
+    
+    @property
+    def frame_shape(self):
+        ''' Return the shape of the video frames. '''
+        return self.container.streams.video[0].codec_context.height, self.container.streams.video[0].codec_context.width
+
+    @property
+    def nframes(self):
+        ''' Return the number of frames in the video. '''
+        return self.container.streams.video[0].frames
+
+    @property
+    def dtype(self):
+        return np.uint8
+
+    @property
+    def shape(self):
+        if self.read_format == 'gray':
+            return (self.nframes, *self.frame_shape)
+        elif self.read_format in ['rgb24', 'bgr24', 'yuv420p']:
+            return (self.nframes, *self.frame_shape, 3)
+        else:
+            raise NotImplementedError(f"shape not implemented for format {self.read_format}")
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    @property
+    def size(self):
+        return np.product(self.shape)
+    
+    @staticmethod
+    def static_shape(filename):
+        ''' Get the shape of a video (static method). 
+        
+        Args:
+            filename (str): path to the video file'''
+        with av.open(filename) as container:
+            stream = container.streams.video[0]
+            shape = np.array([stream.frames, stream.codec_context.height, stream.codec_context.width])
+        return shape
+
+
+def warn_transcode(msg):
+    warnings.warn(msg, stacklevel=2)
+    warnings.warn('Consider transcoding (ffmpeg -y -i "input.mp4" -c:v libx264 -pix_fmt yuv420p -preset superfast -crf 23 "output.mp4").', stacklevel=2)
+
+
+class VideoWriter:
+    ''' A class to write video files, using PyAV.
+    
+    Args:
+        filename (str): path to the output video file
+        codec (str): codec to use. Default is 'libx264'. Other options include 'h264', etc. Check VideoWriter.codecs_available for a list of available codecs.
+        fps (int): frames per second. Default is 25.
+        frame_shape (tuple): shape of the frames. Default is None (to be determined by the first frame written).
+        **kwargs: additional keyword arguments to pass to the codec. Check e.g. `ffmpeg -h encoder=libx264`
+    '''
+
+    def __init__(self, filename, codec='libx264', fps=25, frame_shape=None, pix_fmt='yuv420p', **kwargs):
+        self.container = av.open(filename, mode='w')
+        #define codec defaults:
+        if codec == 'h264':
+            kwargs = dict(bit_rate=1000000, pix_fmt='yuv420p').update(kwargs)
+        elif codec == 'libx264':
+            kwargs = dict(crf=23, preset='superfast', pix_fmt='yuv420p').update(kwargs)
+        #add video stream, if frame_shape is provided (otherwise create when the first frame is written)
+        if frame_shape is not None:
+            self.stream = self.container.add_stream(codec, rate=fps, height=frame_shape[0], width=frame_shape[1], pix_fmt=pix_fmt, **kwargs)
+        else:
+            self.stream = None
+            self.codec = codec
+            self.fps = fps
+            self.pix_fmt = pix_fmt
+            self.kwargs = kwargs if kwargs is not None else {}
+
+    def write(self, im):
+        """ Write a frame to the video file
+        
+        Args:
+            im (np.ndarray): frame to write
+        """
+        if im.ndim == 2:
+            format = 'gray'
+        elif im.ndim == 3 and im.shape[2] == 3:
+            format = 'rgb24'
+        else:
+            raise ValueError(f"Unsupported frame shape: {im.shape}")
+        if self.stream is None:
+            self.stream = self.container.add_stream(self.codec, rate=self.fps, height=im.shape[0], width=im.shape[1], pix_fmt=self.pix_fmt, **self.kwargs)
+        out_frame = av.VideoFrame.from_ndarray(im, format=format)
+        for packet in self.stream.encode(out_frame):
+            self.container.mux(packet)
+
+    def write_frames(self, frames):
+        for frame in frames:
+            self.write(frame)
+
+    def write_audio(self, audio_chunk, audio_rate, pts=None, stream_args=None):
+        """ Write an audio chunk to the video file
+
+        Args:
+            audio_chunk (np.ndarray): audio chunk to write
+            audio_rate (int): audio sample rate
+            pts (int): presentation timestamp of the audio chunk. Default is None (auto-increment)
+            stream_args (dict): additional arguments to pass to the audio stream.
+        """
+        if stream_args is None:
+            stream_args = dict(codec_name='aac', rate=44100)
+        if self.audio_stream is None:
+            self.audio_stream = self.container.add_stream(**stream_args)
+            self.channel_layout = 'stereo' if (audio_chunk.ndim > 1 and audio_chunk.shape[0] == 2) else 'mono'
+            self.audio_pts = 0
+        audio_frame = av.AudioFrame.from_ndarray(audio_chunk, layout=self.channel_layout)
+        audio_frame.sample_rate = audio_rate
+        audio_frame.pts = self.audio_pts if pts is None else pts
+        self.audio_pts = audio_frame.pts + audio_chunk.shape[1]
+        for packet in self.audio_stream.encode(audio_frame):
+            self.container.mux(packet)
+
+    def close(self):
+        """ Close the video file. """
+        try:
+            for packet in self.stream.encode():
+                self.container.mux(packet)
+        except:
+            pass
+        
+        try: 
+            for packet in self.audio_stream.encode():
+                self.container.mux(packet)
+        except:
+            pass
+
+        self.container.close()
+
+    def __del__(self):
+        self.close()
+            
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False
+    
+    @property
+    def codecs_available(self):
+        ''' Return a list of available codecs. '''
+        return sorted(av.codec.codecs_available)

daio-0.0.1.dist-info/LICENSE

@@ -0,0 +1,22 @@
+
+The MIT License (MIT)
+
+Copyright (c) jlab.berlin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

daio-0.0.1.dist-info/METADATA

@@ -0,0 +1,39 @@
+Metadata-Version: 2.1
+Name: daio
+Version: 0.0.1
+Summary: video and data IO tools for Python
+Author: jlab.berlin
+License: MIT
+Project-URL: Homepage, https://github.com/danionella/daio
+Keywords: python,HDF5,data IO,video IO
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: tqdm
+Requires-Dist: av
+Requires-Dist: h5py
+Requires-Dist: hdf5plugin
+
+![Python Version](https://img.shields.io/badge/python-3.8+-blue)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+![tests](https://github.com/danionella/daio/actions/workflows/test.yml/badge.svg)
+[![PyPI - Version](https://img.shields.io/pypi/v/daio)](https://pypi.org/project/daio/)
+[![Conda Version](https://img.shields.io/conda/v/danionella/daio)](https://anaconda.org/danionella/daio)
+![GitHub last commit](https://img.shields.io/github/last-commit/danionella/daio)
+
+# daio
+Video and data IO tools for Python.
+
+Links: [API documentation](http://danionella.github.io/my_package), [GitHub repository](https://github.com/danionella/my_package)
+
+## Installation
+- via conda or mamba: `conda install danionella::daio `
+- if you prefer pip: `pip install daio`
+- for development, clone this repository change to the directory containing `pyproject.toml` 
+    - `conda env create -n env_name -f environment.yml`
+    - `conda activate env_name`
+    - `pip install -e .`
+
+## Use 
+TODO.

daio-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+daio/h5.py,sha256=IjKugDGb6QzHqmGwj_7QSvV1ZbVFf5zJqvxLVr3Rf0I,7305
+daio/video.py,sha256=24_0q10WM5TEfu7em0d7C2XaxyAwmbzhQ8jz6BR1ZIo,10862
+daio-0.0.1.dist-info/LICENSE,sha256=mUNAxIWH48FIpcbQFXkPZWqds0NLZZedQaVjySqcCak,1074
+daio-0.0.1.dist-info/METADATA,sha256=CJODXFFrA-NddeJ9dNGbAEPsEE16-GapVgrtXQDW-lE,1452
+daio-0.0.1.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+daio-0.0.1.dist-info/top_level.txt,sha256=SysaGdJINFYG21hBxF3d5zqIDJXhPHL0lc65lYd1N04,5
+daio-0.0.1.dist-info/RECORD,,

daio-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.6.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

daio-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+daio