flaxdiff 0.1.38__py3-none-any.whl → 0.2.0__py3-none-any.whl

flaxdiff/data/sources/base.py

@@ -0,0 +1,129 @@
+from abc import ABC, abstractmethod
+import grain.python as pygrain
+from typing import Dict, Any, Callable, List, Optional
+import jax.numpy as jnp
+from functools import partial
+
+
+class DataSource(ABC):
+    """Base class for all data sources in FlaxDiff."""
+    
+    @abstractmethod
+    def get_source(self, path_override: str) -> Any:
+        """Return the data source object.
+        
+        Args:
+            path_override: Path to the dataset, overriding the default.
+            
+        Returns:
+            A data source object compatible with grain or other loaders.
+        """
+        pass
+    
+    @staticmethod
+    def create(source_type: str, **kwargs) -> 'DataSource':
+        """Factory method to create a data source of the specified type.
+        
+        Args:
+            source_type: Type of the data source ("image", "video", etc.)
+            **kwargs: Additional arguments for the specific data source.
+            
+        Returns:
+            An instance of a DataSource subclass.
+        """
+        from .images import ImageTFDSSource, ImageGCSSource, CombinedImageGCSSource
+        from .videos import VideoTFDSSource, VideoLocalSource
+        
+        source_map = {
+            "image_tfds": ImageTFDSSource,
+            "image_gcs": ImageGCSSource,
+            "image_combined_gcs": CombinedImageGCSSource,
+            "video_tfds": VideoTFDSSource,
+            "video_local": VideoLocalSource
+        }
+        
+        if source_type not in source_map:
+            raise ValueError(f"Unknown source type: {source_type}")
+        return source_map[source_type](**kwargs)
+
+
+class DataAugmenter(ABC):
+    """Base class for all data augmenters in FlaxDiff."""
+    
+    @abstractmethod
+    def create_transform(self, **kwargs) -> Callable[[], pygrain.MapTransform]:
+        """Create a transformation function for the data.
+        
+        Args:
+            **kwargs: Additional arguments for the transformation.
+            
+        Returns:
+            A callable that returns a pygrain.MapTransform instance.
+        """
+        pass
+    
+    @staticmethod
+    def create(augmenter_type: str, **kwargs) -> 'DataAugmenter':
+        """Factory method to create a data augmenter of the specified type.
+        
+        Args:
+            augmenter_type: Type of the data augmenter ("image", "video", etc.)
+            **kwargs: Additional arguments for the specific augmenter.
+            
+        Returns:
+            An instance of a DataAugmenter subclass.
+        """
+        from .images import ImageTFDSAugmenter, ImageGCSAugmenter
+        from .videos import VideoAugmenter
+        
+        augmenter_map = {
+            "image_tfds": ImageTFDSAugmenter,
+            "image_gcs": ImageGCSAugmenter,
+            "video": VideoAugmenter
+        }
+        
+        if augmenter_type not in augmenter_map:
+            raise ValueError(f"Unknown augmenter type: {augmenter_type}")
+        
+        return augmenter_map[augmenter_type](**kwargs)
+
+
+class MediaDataset:
+    """A class combining a data source and an augmenter for a complete dataset."""
+    
+    def __init__(self, 
+                 source: DataSource, 
+                 augmenter: DataAugmenter,
+                 media_type: str = "image"):
+        """Initialize a MediaDataset.
+        
+        Args:
+            source: The data source.
+            augmenter: The data augmenter.
+            media_type: Type of media ("image", "video", etc.)
+        """
+        self.source = source
+        self.augmenter = augmenter
+        self.media_type = media_type
+    
+    def get_source(self, path_override: str) -> Any:
+        """Get the data source.
+        
+        Args:
+            path_override: Path to override the default data source path.
+            
+        Returns:
+            A data source object.
+        """
+        return self.source.get_source(path_override)
+    
+    def get_augmenter(self, **kwargs) -> Callable[[], pygrain.MapTransform]:
+        """Get the augmenter transformation.
+        
+        Args:
+            **kwargs: Additional arguments for the augmenter.
+            
+        Returns:
+            A callable that returns a pygrain.MapTransform instance.
+        """
+        return self.augmenter.create_transform(**kwargs)

flaxdiff/data/sources/images.py

@@ -0,0 +1,309 @@
+import cv2
+import jax.numpy as jnp
+import grain.python as pygrain
+from flaxdiff.utils import AutoTextTokenizer
+from typing import Dict, Any, Callable, List, Optional
+import random
+import augmax
+import jax
+import os
+import struct as st
+from functools import partial
+import numpy as np
+from .base import DataSource, DataAugmenter
+
+
+# ----------------------------------------------------------------------------------
+# Utility functions
+# ----------------------------------------------------------------------------------
+
+def unpack_dict_of_byte_arrays(packed_data):
+    """Unpacks a dictionary of byte arrays from a packed binary format."""
+    unpacked_dict = {}
+    offset = 0
+    while offset < len(packed_data):
+        # Unpack the key length
+        key_length = st.unpack_from('I', packed_data, offset)[0]
+        offset += st.calcsize('I')
+        # Unpack the key bytes and convert to string
+        key = packed_data[offset:offset+key_length].decode('utf-8')
+        offset += key_length
+        # Unpack the byte array length
+        byte_array_length = st.unpack_from('I', packed_data, offset)[0]
+        offset += st.calcsize('I')
+        # Unpack the byte array
+        byte_array = packed_data[offset:offset+byte_array_length]
+        offset += byte_array_length
+        unpacked_dict[key] = byte_array
+    return unpacked_dict
+
+
+# ----------------------------------------------------------------------------------
+# Image augmentation utilities
+# ----------------------------------------------------------------------------------
+
+def image_augmenter(image, image_scale, method=cv2.INTER_AREA):
+    """Basic image augmentation: convert color and resize."""
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image = cv2.resize(image, (image_scale, image_scale),
+                       interpolation=method)
+    return image
+
+
+PROMPT_TEMPLATES = [
+    "a photo of a {}",
+    "a photo of a {} flower",
+    "This is a photo of a {}",
+    "This is a photo of a {} flower",
+    "A photo of a {} flower",
+]
+
+
+def labelizer_oxford_flowers102(path):
+    """Creates a label generator for Oxford Flowers 102 dataset."""
+    with open(path, "r") as f:
+        textlabels = [i.strip() for i in f.readlines()]
+
+    def load_labels(sample):
+        raw = textlabels[int(sample['label'])]
+        # randomly select a prompt template
+        template = random.choice(PROMPT_TEMPLATES)
+        # format the template with the label
+        caption = template.format(raw)
+        # return the caption
+        return caption
+    return load_labels
+
+
+# ----------------------------------------------------------------------------------
+# TFDS Image Source
+# ----------------------------------------------------------------------------------
+
+class ImageTFDSSource(DataSource):
+    """Data source for TensorFlow Datasets (TFDS) image datasets."""
+    
+    def __init__(self, name: str, use_tf: bool = True, split: str = "all"):
+        """Initialize a TFDS image data source.
+        
+        Args:
+            name: Name of the TFDS dataset.
+            use_tf: Whether to use TensorFlow for loading.
+            split: Dataset split to use.
+        """
+        self.name = name
+        self.use_tf = use_tf
+        self.split = split
+    
+    def get_source(self, path_override: str) -> Any:
+        """Get the TFDS data source.
+        
+        Args:
+            path_override: Override path for the dataset.
+            
+        Returns:
+            A TFDS dataset.
+        """
+        import tensorflow_datasets as tfds
+        if self.use_tf:
+            return tfds.load(self.name, split=self.split, shuffle_files=True)
+        else:
+            return tfds.data_source(self.name, split=self.split, try_gcs=False)
+
+
+class ImageTFDSAugmenter(DataAugmenter):
+    """Augmenter for TFDS image datasets."""
+    
+    def __init__(self, label_path: str = "/home/mrwhite0racle/tensorflow_datasets/oxford_flowers102/2.1.1/label.labels.txt"):
+        """Initialize a TFDS image augmenter.
+        
+        Args:
+            label_path: Path to the labels file for datasets like Oxford Flowers.
+        """
+        self.label_path = label_path
+    
+    def create_transform(self, image_scale: int = 256, method: Any = None) -> Callable[[], pygrain.MapTransform]:
+        """Create a transform for TFDS image datasets.
+        
+        Args:
+            image_scale: Size to scale images to.
+            method: Interpolation method for resizing.
+            
+        Returns:
+            A callable that returns a pygrain.MapTransform.
+        """
+        labelizer = labelizer_oxford_flowers102(self.label_path)
+        
+        if image_scale > 256:
+            interpolation = cv2.INTER_CUBIC
+        else:
+            interpolation = cv2.INTER_AREA
+        
+        from torchvision.transforms import v2
+        augments = v2.Compose([
+            v2.RandomHorizontalFlip(p=0.5),
+            v2.ColorJitter(brightness=0.2, contrast=0.05, saturation=0.2)
+        ])
+        
+        class TFDSTransform(pygrain.MapTransform):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.tokenize = AutoTextTokenizer(tensor_type="np")
+                
+            def map(self, element) -> Dict[str, jnp.array]:
+                image = element['image']
+                image = cv2.resize(image, (image_scale, image_scale),
+                                  interpolation=interpolation)
+                image = augments(image)
+                
+                caption = labelizer(element)
+                results = self.tokenize(caption)
+                return {
+                    "image": image,
+                    "text": {
+                        "input_ids": results['input_ids'][0],
+                        "attention_mask": results['attention_mask'][0],
+                    }
+                }
+        
+        return TFDSTransform
+
+
+# ----------------------------------------------------------------------------------
+# GCS Image Source
+# ----------------------------------------------------------------------------------
+
+class ImageGCSSource(DataSource):
+    """Data source for Google Cloud Storage (GCS) image datasets."""
+    
+    def __init__(self, source: str = 'arrayrecord/laion-aesthetics-12m+mscoco-2017'):
+        """Initialize a GCS image data source.
+        
+        Args:
+            source: Path to the GCS dataset.
+        """
+        self.source = source
+    
+    def get_source(self, path_override: str = "/home/mrwhite0racle/gcs_mount") -> Any:
+        """Get the GCS data source.
+        
+        Args:
+            path_override: Base path for GCS mounts.
+            
+        Returns:
+            A grain ArrayRecordDataSource.
+        """
+        records_path = os.path.join(path_override, self.source)
+        records = [os.path.join(records_path, i) for i in os.listdir(
+            records_path) if 'array_record' in i]
+        return pygrain.ArrayRecordDataSource(records)
+
+
+class CombinedImageGCSSource(DataSource):
+    """Data source that combines multiple GCS image datasets."""
+    
+    def __init__(self, sources: List[str] = []):
+        """Initialize a combined GCS image data source.
+        
+        Args:
+            sources: List of paths to GCS datasets.
+        """
+        self.sources = sources
+    
+    def get_source(self, path_override: str = "/home/mrwhite0racle/gcs_mount") -> Any:
+        """Get the combined GCS data source.
+        
+        Args:
+            path_override: Base path for GCS mounts.
+            
+        Returns:
+            A grain ArrayRecordDataSource.
+        """
+        records_paths = [os.path.join(path_override, source) for source in self.sources]
+        records = []
+        for records_path in records_paths:
+            records += [os.path.join(records_path, i) for i in os.listdir(
+                records_path) if 'array_record' in i]
+        return pygrain.ArrayRecordDataSource(records)
+
+
+class ImageGCSAugmenter(DataAugmenter):
+    """Augmenter for GCS image datasets."""
+    
+    def __init__(self, labelizer: Callable = None):
+        """Initialize a GCS image augmenter.
+        
+        Args:
+            labelizer: Function to extract text labels from samples.
+        """
+        self.labelizer = labelizer or (lambda sample: sample['txt'])
+    
+    def create_transform(self, image_scale: int = 256, method: Any = None) -> Callable[[], pygrain.MapTransform]:
+        """Create a transform for GCS image datasets.
+        
+        Args:
+            image_scale: Size to scale images to.
+            method: Interpolation method for resizing.
+            
+        Returns:
+            A callable that returns a pygrain.MapTransform.
+        """
+        labelizer = self.labelizer
+        
+        class GCSTransform(pygrain.MapTransform):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.auto_tokenize = AutoTextTokenizer(tensor_type="np")
+                self.image_augmenter = partial(image_augmenter, image_scale=image_scale, method=method)
+                
+            def map(self, element) -> Dict[str, jnp.array]:
+                element = unpack_dict_of_byte_arrays(element)
+                image = np.asarray(bytearray(element['jpg']), dtype="uint8")
+                image = cv2.imdecode(image, cv2.IMREAD_UNCHANGED)
+                image = self.image_augmenter(image)
+                caption = labelizer(element).decode('utf-8')
+                results = self.auto_tokenize(caption)
+                return {
+                    "image": image,
+                    "text": {
+                        "input_ids": results['input_ids'][0],
+                        "attention_mask": results['attention_mask'][0],
+                    }
+                }
+        
+        return GCSTransform
+
+
+# ----------------------------------------------------------------------------------
+# Legacy compatibility functions
+# ----------------------------------------------------------------------------------
+
+# These functions maintain backward compatibility with existing code
+
+def data_source_tfds(name, use_tf=True, split="all"):
+    """Legacy function for TFDS data sources."""
+    source = ImageTFDSSource(name=name, use_tf=use_tf, split=split)
+    return source.get_source
+
+
+def tfds_augmenters(image_scale, method):
+    """Legacy function for TFDS augmenters."""
+    augmenter = ImageTFDSAugmenter()
+    return augmenter.create_transform(image_scale=image_scale, method=method)
+
+
+def data_source_gcs(source='arrayrecord/laion-aesthetics-12m+mscoco-2017'):
+    """Legacy function for GCS data sources."""
+    source_obj = ImageGCSSource(source=source)
+    return source_obj.get_source
+
+
+def data_source_combined_gcs(sources=[]):
+    """Legacy function for combined GCS data sources."""
+    source_obj = CombinedImageGCSSource(sources=sources)
+    return source_obj.get_source
+
+
+def gcs_augmenters(image_scale, method):
+    """Legacy function for GCS augmenters."""
+    augmenter = ImageGCSAugmenter()
+    return augmenter.create_transform(image_scale=image_scale, method=method)

flaxdiff/data/sources/utils.py

@@ -0,0 +1,158 @@
+
+import numpy as np
+from decord.video_reader import VideoReader
+from decord.audio_reader import AudioReader
+
+from decord.ndarray import cpu
+from decord import ndarray as _nd
+from decord.bridge import bridge_out
+
+class AVReader(object):
+    """Individual audio video reader with convenient indexing function.
+
+    Parameters
+    ----------
+    uri: str
+        Path of file.
+    ctx: decord.Context
+        The context to decode the file, can be decord.cpu() or decord.gpu().
+    sample_rate: int, default is -1
+        Desired output sample rate of the audio, unchanged if `-1` is specified.
+    mono: bool, default is True
+        Desired output channel layout of the audio. `True` is mono layout. `False` is unchanged.
+    width : int, default is -1
+        Desired output width of the video, unchanged if `-1` is specified.
+    height : int, default is -1
+        Desired output height of the video, unchanged if `-1` is specified.
+    num_threads : int, default is 0
+        Number of decoding thread, auto if `0` is specified.
+    fault_tol : int, default is -1
+        The threshold of corupted and recovered frames. This is to prevent silent fault
+        tolerance when for example 50% frames of a video cannot be decoded and duplicate
+        frames are returned. You may find the fault tolerant feature sweet in many cases,
+        but not for training models. Say `N = # recovered frames`
+        If `fault_tol` < 0, nothing will happen.
+        If 0 < `fault_tol` < 1.0, if N > `fault_tol * len(video)`, raise `DECORDLimitReachedError`.
+        If 1 < `fault_tol`, if N > `fault_tol`, raise `DECORDLimitReachedError`.
+    """
+
+    def __init__(
+        self, uri, ctx=cpu(0), sample_rate=-1, mono=True, width=-1, height=-1, num_threads=0, fault_tol=-1
+    ):
+        self.__audio_reader = AudioReader(uri, ctx, sample_rate, mono)
+        self.__audio_reader.add_padding()
+        if hasattr(uri, "read"):
+            uri.seek(0)
+        self.__video_reader = VideoReader(uri, ctx, width, height, num_threads, fault_tol)
+        self.__video_reader.seek(0)
+        
+    def __del__(self):
+        del self.__video_reader
+        del self.__audio_reader
+
+    def __len__(self):
+        """Get length of the video. Note that sometimes FFMPEG reports inaccurate number of frames,
+        we always follow what FFMPEG reports.
+        Returns
+        -------
+        int
+            The number of frames in the video file.
+        """
+        return len(self.__video_reader)
+
+    def __getitem__(self, idx):
+        """Get audio samples and video frame at `idx`.
+
+        Parameters
+        ----------
+        idx : int or slice
+            The frame index, can be negative which means it will index backwards,
+            or slice of frame indices.
+
+        Returns
+        -------
+        (ndarray/list of ndarray, ndarray)
+            First element is samples of shape CxS or a list of length N containing samples of shape CxS,
+            where N is the number of frames, C is the number of channels,
+            S is the number of samples of the corresponding frame.
+
+            Second element is Frame of shape HxWx3 or batch of image frames with shape NxHxWx3,
+            where N is the length of the slice.
+        """
+        assert self.__video_reader is not None and self.__audio_reader is not None
+        if isinstance(idx, slice):
+            return self.get_batch(range(*idx.indices(len(self.__video_reader))))
+        if idx < 0:
+            idx += len(self.__video_reader)
+        if idx >= len(self.__video_reader) or idx < 0:
+            raise IndexError("Index: {} out of bound: {}".format(idx, len(self.__video_reader)))
+        audio_start_idx, audio_end_idx = self.__video_reader.get_frame_timestamp(idx)
+        audio_start_idx = self.__audio_reader._time_to_sample(audio_start_idx)
+        audio_end_idx = self.__audio_reader._time_to_sample(audio_end_idx)
+        results = (self.__audio_reader[audio_start_idx:audio_end_idx], self.__video_reader[idx])
+        self.__video_reader.seek(0)
+        return results
+
+    def get_batch(self, indices):
+        """Get entire batch of audio samples and video frames.
+
+        Parameters
+        ----------
+        indices : list of integers
+            A list of frame indices. If negative indices detected, the indices will be indexed from backward
+        Returns
+        -------
+        (list of ndarray, ndarray)
+            First element is a list of length N containing samples of shape CxS,
+            where N is the number of frames, C is the number of channels,
+            S is the number of samples of the corresponding frame.
+
+            Second element is Frame of shape HxWx3 or batch of image frames with shape NxHxWx3,
+            where N is the length of the slice.
+
+        """
+        assert self.__video_reader is not None and self.__audio_reader is not None
+        indices = self._validate_indices(indices)
+        audio_arr = []
+        prev_video_idx = None
+        prev_audio_end_idx = None
+        for idx in list(indices):
+            frame_start_time, frame_end_time = self.__video_reader.get_frame_timestamp(idx)
+            # timestamp and sample conversion could have some error that could cause non-continuous audio
+            # we detect if retrieving continuous frame and make the audio continuous
+            if prev_video_idx and idx == prev_video_idx + 1:
+                audio_start_idx = prev_audio_end_idx
+            else:
+                audio_start_idx = self.__audio_reader._time_to_sample(frame_start_time)
+            audio_end_idx = self.__audio_reader._time_to_sample(frame_end_time)
+            audio_arr.append(self.__audio_reader[audio_start_idx:audio_end_idx])
+            prev_video_idx = idx
+            prev_audio_end_idx = audio_end_idx
+        results = (audio_arr, self.__video_reader.get_batch(indices))
+        self.__video_reader.seek(0)
+        return results
+
+    def _get_slice(self, sl):
+        audio_arr = np.empty(shape=(self.__audio_reader.shape()[0], 0), dtype="float32")
+        for idx in list(sl):
+            audio_start_idx, audio_end_idx = self.__video_reader.get_frame_timestamp(idx)
+            audio_start_idx = self.__audio_reader._time_to_sample(audio_start_idx)
+            audio_end_idx = self.__audio_reader._time_to_sample(audio_end_idx)
+            audio_arr = np.concatenate(
+                (audio_arr, self.__audio_reader[audio_start_idx:audio_end_idx].asnumpy()), axis=1
+            )
+        results = (bridge_out(_nd.array(audio_arr)), self.__video_reader.get_batch(sl))
+        self.__video_reader.seek(0)
+        return results
+
+    def _validate_indices(self, indices):
+        """Validate int64 integers and convert negative integers to positive by backward search"""
+        assert self.__video_reader is not None and self.__audio_reader is not None
+        indices = np.array(indices, dtype=np.int64)
+        # process negative indices
+        indices[indices < 0] += len(self.__video_reader)
+        if not (indices >= 0).all():
+            raise IndexError("Invalid negative indices: {}".format(indices[indices < 0] + len(self.__video_reader)))
+        if not (indices < len(self.__video_reader)).all():
+            raise IndexError("Out of bound indices: {}".format(indices[indices >= len(self.__video_reader)]))
+        return indices