pyremotedata.dataloader — PyRemoteData 0.0.34 documentation

Skip to content

# Source code for pyremotedata.dataloader

    ## PSA: This script is not ready for general use, as it still hardcoded to the specific use case of the project it was developed for.

    # Standard library imports
    import os
    import time
    import warnings
    from queue import Queue
    from threading import Thread
    from typing import Union, Optional, Tuple, List, Callable

    # Dependency imports
    import torch
    from torch.utils.data import DataLoader, IterableDataset#, TensorDataset
    from torchvision.io import read_image
    from torchvision.io.image import ImageReadMode

    # Backend imports
    from .implicit_mount import RemotePathIterator

    class RemotePathDataset(IterableDataset):
        '''
        Creates a :py:class:`torch.utils.data.IterableDataset` from a :py:class:`pyremotedata.implicit_mount.RemotePathIterator`.

        By default the dataset will return the image as a tensor and the remote path as a string.

        **Hierarchical mode**

        If `hierarchical` >= 1, the dataset is in "Hierarchical mode" and will return the image as a tensor and the label as a list of integers (class indices for each level in the hierarchy).

        The `class_handles` property can be used to get the class-idx mappings for the dataset.

        By default the dataset will use a parser which assumes that the hierarchical levels are encoded in the remote path as directories like so:

        `.../level_n/.../level_1/level_0/image.jpg`

        Where `n = (hierarchical - 1)` and `level_0` is the leaf level.

        Args:
            remote_path_iterator (RemotePathIterator): The :py:class:`pyremotedata.implicit_mount.RemotePathIterator` to create the dataset from.
            prefetch (int): The number of items to prefetch from the :py:class:`pyremotedata.implicit_mount.RemotePathIterator`.
            transform (callable, optional): A function/transform that takes in an image as a :py:class:`torch.Tensor` and returns a transformed version.
            target_transform (callable, optional): A function/transform that takes in the label (after potential parsing by `parse_hierarchical`) and transforms it.
            device (torch.device, optional): The device to move the tensors to.
            dtype (torch.dtype, optional): The data type to convert the tensors to.
            hierarchical (int, optional): The number of hierarchical levels to use for the labels. Default: 0, i.e. no hierarchy.
            hierarchy_parser (callable, optional): A function to parse the hierarchical levels from the remote path. Default: None, i.e. use the default parser.
            return_remote_path (bool, optional): Whether to return the remote path. Default: False.
            return_local_path (bool, optional): Whether to return the local path. Default: False.
            verbose (bool, optional): Whether to print verbose output. Default: False.

        Yields:
            (tuple): A tuple containing the following elements:
                - (torch.Tensor): The image as a tensor.
                - (Union[str, List[int]]): The label as the remote path or as a list of class indices.
                - (Optional[str]): The local path, if `return_local_path` is True.
                - (Optional[str]): The remote path, if `return_remote_path` is True.
        '''
        def __init__(
                self,
                remote_path_iterator : "RemotePathIterator",
                prefetch: int=64,
                transform : Optional[Callable]=None,
                target_transform : Optional[Callable]=None,
                device: Union["torch.device", None]=None,
                dtype: Union[torch.dtype, None]=None,
                hierarchical: int=0,
                hierarchy_parser: Optional[Callable]=None,
                shuffle: bool=False,
                return_remote_path: bool=False,
                return_local_path: bool=False,
                verbose: bool=False
            ):
            # Check if remote_path_iterator is of type RemotePathIterator
            if not isinstance(remote_path_iterator, RemotePathIterator):
                raise ValueError("Argument remote_path_iterator must be of type pyremotedata.implicit_mount.RemotePathIterator.")
            # Check if prefetch is an integer
            if not isinstance(prefetch, int):
                raise ValueError("Argument prefetch must be an integer.")
            # Check if prefetch is greater than 0
            if prefetch < 1:
                raise ValueError("Argument prefetch must be greater than 0.")

            ## General parameters
            assert isinstance(verbose, bool), ValueError("Argument verbose must be a boolean.")
            self.verbose : bool = verbose

            ## PyTorch specific parameters
            # Get the classes and their indices
            self.hierarchical : int = hierarchical
            if self.hierarchical < 1:
                self.hierarchical = False
                def error_hierarchy(*args, **kwargs):
                    raise ValueError("Hierarchical mode disabled (`hierarchical` < 1), but `pyremotedata.dataloader.RemotePathDataset.parse_hierarchy` function was called.")
                self.parse_hierarchy = error_hierarchy
            else:
                if hierarchy_parser is None:
                    self.parse_hierarchy = lambda path: path.split('/')[-(1 + self.hierarchical):-1]
                elif callable(hierarchy_parser):
                    self.parse_hierarchy = hierarchy_parser
                else:
                    raise ValueError("Argument `hierarchy_parser` must be a callable or None.")
                self.classes = [[] for _ in range(self.hierarchical)]
                self.n_classes = [0 for _ in range(self.hierarchical)]
                self.class_to_idx = [{} for _ in range(self.hierarchical)]
                self.idx_to_class = [{} for _ in range(self.hierarchical)]
                for path in remote_path_iterator.remote_paths:
                    for level, cls in enumerate(reversed(self.parse_hierarchy(path))):
                        if cls in self.classes[level]:
                            continue
                        if level >= self.hierarchical:
                            raise ValueError(f"Error parsing class from {path}. Got {cls} at level {level}, but the number of specified hierarchical levels is {self.hierarchical} with levels 0-{self.hierarchical-1}.")
                        self.classes[level].append(cls)
                for level in range(self.hierarchical):
                    # self.classes[level] = sorted(list(set([path.split('/')[-2-level] for path in remote_path_iterator.remote_paths])))
                    self.classes[level] = sorted(self.classes[level])
                    self.n_classes[level] = len(self.classes[level])
                    self.class_to_idx[level] = {self.classes[level][i]: i for i in range(len(self.classes[level]))}
                    self.idx_to_class[level] = {i: self.classes[level][i] for i in range(len(self.classes[level]))}

            # Set the transforms
            self.transform = transform
            self.target_transform = target_transform

            # Set the device and dtype
            self.device = device
            self.dtype = dtype

            ## Backend specific parameters
            # Store the remote_path_iterator backend
            self.remote_path_iterator = remote_path_iterator
            self.return_remote_path = return_remote_path
            self.return_local_path = return_local_path
            self.shuffle = shuffle

            ## Multi-threading parameters
            # Set the number of workers (threads) for (multi)processing
            self.num_workers = 1

            # We don't want to start the buffer filling thread until the dataloader is called for iteration
            self.producer_thread = None
            # We need to keep track of whether the buffer filling thread has been initiated or not
            self.thread_initiated = False

            # Initialize the worker threads
            self.consumer_threads = []
            self.consumers = 0
            self.stop_consumer_threads = True

            # Set the buffer filling parameters (Watermark Buffering)
            self.buffer_minfill, self.buffer_maxfill = 0.4, 0.6

            # Initialize the buffers
            self.buffer = Queue(maxsize=prefetch)  # Tune maxsize as needed
            self.processed_buffer = Queue(maxsize=prefetch)  # Tune maxsize as needed

        @property
        def class_handles(self):
            return {
                'classes': self.classes,
                'n_classes': self.n_classes,
                'class_to_idx': self.class_to_idx,
                'idx_to_class': self.idx_to_class,
                'hierarchical': self.hierarchical
            }

        @class_handles.setter
        def class_handles(self, value : dict):
            if not isinstance(value, dict):
                raise ValueError("Argument value must be a dictionary.")
            if value["hierarchical"]:
                assert isinstance(value['classes'], list), ValueError("Argument value['classes'] must be a list, when hierarchical is True.")
            else:
                assert not isinstance(value['classes'], list), ValueError("Argument value['classes'] must not be a list, when hierarchical is False.")
            self.classes = value['classes']
            self.n_classes = value['n_classes']
            self.class_to_idx = value['class_to_idx']
            self.idx_to_class = value['idx_to_class']
            self.hierarchical = value['hierarchical']

        def _shuffle(self):
            if not self.shuffle:
                raise RuntimeError("Shuffle called, but shuffle is set to False.")
            if self.thread_initiated:
                raise RuntimeError("Shuffle called, but buffer filling thread is still active.")
            self.remote_path_iterator.shuffle()

        def _shutdown_and_reset(self):
            # Handle shutdown logic (Should probably be moved to a dedicated reset function, that is called on StopIteration instead or perhaps in __iter__)
            self.stop_consumer_threads = True # Signal the consumer threads to stop
            for i, consumer in enumerate(self.consumer_threads):
                if consumer is None:
                    continue
                if not consumer.is_alive():
                    continue
                print(f"Waiting for worker {i} to finish.")
                consumer.join(timeout = 1 / self.num_workers) # Wait for the consumer thread to finish
            if self.producer_thread is not None:
                self.producer_thread.join(timeout=1) # Wait for the producer thread to finish
                self.producer_thread = None # Reset the producer thread
            self.consumer_threads = [] # Reset the consumer threads
            self.consumers = 0 # Reset the number of consumers
            self.thread_initiated = False # Reset the thread initiated flag
            self.buffer.queue.clear() # Clear the buffer
            self.processed_buffer.queue.clear() # Clear the processed buffer
            self.remote_path_iterator.__del__(force=True)  # Close the remote_path_iterator and clean the temporary directory
            assert self.buffer.qsize() == 0, RuntimeError("Buffer not empty after iterator end.")
            assert self.processed_buffer.qsize() == 0, RuntimeError("Processed buffer not empty after iterator end.")

        def _init_buffer(self):
            # Check if the buffer filling thread has been initiated
            if not self.thread_initiated:
                # Start the buffer filling thread
                self.producer_thread = Thread(target=self._fill_buffer)
                self.producer_thread.daemon = True
                self.producer_thread.start()
                # Set the flag to indicate that the thread has been initiated
                self.thread_initiated = True
            else:
                # Raise an error if the buffer filling thread has already been initiated
                raise RuntimeError("Buffer filling thread already initiated.")

        def _fill_buffer(self):
            # Calculate the min and max fill values for the buffer
            min_fill = int(self.buffer.maxsize * self.buffer_minfill)
            max_fill = int(self.buffer.maxsize * self.buffer_maxfill)

            if self.verbose:
                print(f"Buffer min fill: {min_fill}, max fill: {max_fill}")
                print("Producer thread started.")

            wait_for_min_fill = False
            for item in self.remote_path_iterator:
                # Get the current size of the buffer
                current_buffer_size = self.buffer.qsize()

                # Decide whether to fill the buffer based on its current size
                if not wait_for_min_fill:
                    wait_for_min_fill = (current_buffer_size >= max_fill)

                # Sleep logic which ensures that the buffer doesn't switch between filling and not filling too often (Watermark Buffering)
                while wait_for_min_fill:
                    time.sleep(0.1)
                    current_buffer_size = self.buffer.qsize() # Update the current buffer size
                    wait_for_min_fill = current_buffer_size >= min_fill # Wait until the buffer drops below min_fill

                # Fill the buffer
                self.buffer.put(item)
            if self.verbose:
                print("Producer signalling end of iterator.")
            # Signal the end of the iterator to the consumers by putting None in the buffer until all consumer threads have finished
            while self.consumers > 0:
                time.sleep(0.01)
                self.buffer.put(None)
            if self.verbose:
                print("Producer emptying buffer.")
            # Wait for the consumer threads to finish then clear the buffer
            while self.buffer.qsize() > 0:
                self.buffer.get()
            self.processed_buffer.put(None) # Signal the end of the processed buffer to the main thread
            if self.verbose:
                print("Producer thread finished.")

        def _process_buffer(self):
            if self.verbose:
                print("Consumer thread started.")
            self.consumers += 1
            while True:
                qsize = self.buffer.qsize()
                while qsize < (self.num_workers * 2):
                    time.sleep(0.05 / self.num_workers)
                    qsize = self.buffer.qsize()
                    if self.producer_thread is None or not self.producer_thread.is_alive():
                        break
                # Get the next item from the buffer
                item = self.buffer.get()
                if self.verbose:
                    print("Consumer thread got item")
                # Check if the buffer is empty, signaling the end of the iterator
                if item is None or self.stop_consumer_threads:
                    break  # Close the thread

                # Preprocess the item (e.g. read image, apply transforms, etc.) and put it in the processed buffer
                processed_item = self.parse_item(*item) if item is not None else None
                if processed_item is not None:
                    self.processed_buffer.put(processed_item)
                if self.verbose:
                    print("Consumer thread processed item")

            self.consumers -= 1
            if self.verbose:
                print("Consumer thread finished.")

        def __iter__(self):
            # Check if the buffer filling thread has been initiated
            # If it has, reset the dataloader state and close all threads
            # (Only one iteration is allowed per dataloader instance)
            if self.thread_initiated:
                warnings.warn("Iterator called, but buffer filling thread is still active. Resetting the dataloader state.")
            self._shutdown_and_reset()

            # If shuffle is set to True, shuffle the remote_path_iterator
            if self.shuffle:
                self._shuffle()

            # Initialize the buffer filling thread
            self._init_buffer()

            # Check number of workers
            if self.num_workers == 0:
                self.num_workers = 1
            if self.num_workers < 1:
                raise ValueError("Number of workers must be greater than 0.")

            self.stop_consumer_threads = False

            # Start consumer threads for processing
            for _ in range(self.num_workers):
                consumer_thread = Thread(target=self._process_buffer)
                consumer_thread.daemon = True
                consumer_thread.start()
                self.consumer_threads.append(consumer_thread)

            return self

        def __next__(self):
            # Fetch from processed_buffer instead
            processed_item = self.processed_buffer.get()

            # Check if the processed buffer is empty, signaling the end of the iterator
            if processed_item is None:
                self._shutdown_and_reset()
                raise StopIteration

            # Restart crashed consumer threads
            for thread in self.consumer_threads:
                if not thread.is_alive():
                    self.consumer_threads.remove(thread)
                    new_consumer_thread = Thread(target=self._process_buffer)
                    new_consumer_thread.daemon = True
                    self.consumer_threads.append(new_consumer_thread)

            # Otherwise, return the processed item
            return processed_item

        def __len__(self):
            return len(self.remote_path_iterator)

        def parse_item(self, local_path : str, remote_path : str) -> Union[Tuple[torch.Tensor, Union[str, List[int]]], Tuple[torch.Tensor, Union[str, List[int]], str], Tuple[torch.Tensor, Union[str, List[int]], str, str]]:
            ## Image processing
            # Check if image format is supported (jpeg/jpg/png)
            image_type = os.path.splitext(local_path)[-1]
            if image_type not in ['.JPG', '.JPEG', '.PNG', '.jpg', '.jpeg', '.png']:
                # raise ValueError(f"Image format of {remote_path} ({image_type}) is not supported.")
                # Instead of raising an error, we can skip the image instead
                return None
            try:
                image = read_image(local_path, mode=ImageReadMode.RGB)
            except Exception as e:
                print(f"Error reading image {remote_path} ({e}).")
                return None
            # Remove the alpha channel if present
            if image.shape[0] == 4:
                image = image[:3]
            if image.shape[0] != 3:
                print(f"Error reading image {remote_path}.")
                return None
            # Apply transforms (preprocessing)
            if self.transform:
                image = self.transform(image)
            if self.dtype is not None:
                image = image.to(dtype=self.dtype)
            if self.device is not None:
                image = image.to(device=self.device)

            if self.hierarchical:
                ## Label processing
                # Get the label by parsing the remote path
                hierarchy = self.parse_hierarchy(remote_path)
                # Transform the species name to the label index
                label = [self.class_to_idx[level][cls] for level, cls in enumerate(hierarchy)]
                if len(label) == 0:
                    raise ValueError(f"Error parsing label from {remote_path}.")
            else:
                label = remote_path

            # Apply label transforms
            if self.target_transform:
                label = self.target_transform(label)
            # TODO: Does the label need to be converted to a tensor and moved to the device? (Probably not) If so, does it need to be a long tensor?
            # if self.device is not None:
            #     label = label.to(device=self.device)

            ## Return the image and label (and optionally the local and remote paths)
            if self.return_local_path and self.return_remote_path:
                return image, label, local_path, remote_path
            elif self.return_local_path:
                return image, label, local_path
            elif self.return_remote_path:
                return image, label, remote_path
            else:
                return image, label

    class RemotePathDataLoader(DataLoader):
        """
        A custom :py:class:`torch.utils.data.DataLoader` for :py:class:`pyremotedata.dataloader.RemotePathDataset`.

        This DataLoader is designed to work with :py:class:`pyremotedata.dataloader.RemotePathDataset` and does not support all the arguments of the standard :py:class:`torch.utils.data.DataLoader`.

        Unsupported arguments:
        - sampler
        - batch_sampler

        Args:
            dataset (RemotePathDataset): The :py:class:`pyremotedata.dataloader.RemotePathDataset` dataset to load from.
            num_workers (int, optional): The number of worker threads to use for loading. Default: 0. Must be greater than 0.
            shuffle (bool, optional): Whether to shuffle the dataset between epochs. Default: False.
        """
        def __init__(self, dataset: "RemotePathDataset", num_workers : int=0, shuffle : bool=False, *args, **kwargs):
            # Snipe arguments from the user which would break the custom dataloader (e.g. sampler, shuffle, etc.)
            unsupported_kwargs = ['sampler', 'batch_sampler']
            for unzkw in unsupported_kwargs:
                value = kwargs.pop(unzkw, None)
                if value is not None:
                    warnings.warn(f"Argument {unzkw} is not supported in `pyremotedata.dataloader.RemotePathDataLoader`. {unzkw}={value} will be ignored.")

            # Override the num_workers argument handling (default is 0)
            dataset.num_workers = num_workers
            # Override the shuffle argument handling (default is False)
            dataset.shuffle = shuffle

            if not isinstance(dataset, RemotePathDataset):
                raise ValueError("Argument dataset must be of type `pyremotedata.dataloader.RemotePathDataset`.")

            # Initialize the dataloader
            super(RemotePathDataLoader, self).__init__(
                shuffle=False,
                sampler=None,
                batch_sampler=None,
                num_workers=0,
                dataset=dataset,
                *args,
                **kwargs
            )

        # TODO: This cannot be set before calling super().__init__(), so perhaps it can be overridden after initialization instead
        # def __setattr__(self, name, value):
        #     if name in ['batch_sampler', 'sampler']:
        #         raise (f"Changing {name} is not allowed in this custom DataLoader.")
        #     super(RemotePathDataLoader, self).__setattr__(name, value)