|
|
from io import BufferedIOBase |
|
|
from typing import Any, Callable, Iterable, Iterator, Sized, Tuple |
|
|
|
|
|
from torch.utils.data.datapipes._decorator import functional_datapipe |
|
|
from torch.utils.data.datapipes.datapipe import IterDataPipe |
|
|
from torch.utils.data.datapipes.utils.common import _deprecation_warning |
|
|
from torch.utils.data.datapipes.utils.decoder import ( |
|
|
Decoder, |
|
|
basichandlers as decoder_basichandlers, |
|
|
imagehandler as decoder_imagehandler, |
|
|
extension_extract_fn |
|
|
) |
|
|
|
|
|
__all__ = ["RoutedDecoderIterDataPipe", ] |
|
|
|
|
|
|
|
|
@functional_datapipe('routed_decode') |
|
|
class RoutedDecoderIterDataPipe(IterDataPipe[Tuple[str, Any]]): |
|
|
r""" |
|
|
Decodes binary streams from input DataPipe, yields pathname and decoded data |
|
|
in a tuple (functional name: ``routed_decode``). |
|
|
|
|
|
Args: |
|
|
datapipe: Iterable datapipe that provides pathname and binary stream in tuples |
|
|
handlers: Optional user defined decoder handlers. If ``None``, basic and image decoder |
|
|
handlers will be set as default. If multiple handles are provided, the priority |
|
|
order follows the order of handlers (the first handler has the top priority) |
|
|
key_fn: Function for decoder to extract key from pathname to dispatch handlers. |
|
|
Default is set to extract file extension from pathname |
|
|
|
|
|
Note: |
|
|
When ``key_fn`` is specified returning anything other than extension, the default |
|
|
handler will not work and users need to specify custom handler. Custom handler |
|
|
could use regex to determine the eligibility to handle data. |
|
|
""" |
|
|
|
|
|
def __init__(self, |
|
|
datapipe: Iterable[Tuple[str, BufferedIOBase]], |
|
|
*handlers: Callable, |
|
|
key_fn: Callable = extension_extract_fn) -> None: |
|
|
super().__init__() |
|
|
self.datapipe: Iterable[Tuple[str, BufferedIOBase]] = datapipe |
|
|
if not handlers: |
|
|
handlers = (decoder_basichandlers, decoder_imagehandler('torch')) |
|
|
self.decoder = Decoder(*handlers, key_fn=key_fn) |
|
|
_deprecation_warning( |
|
|
type(self).__name__, |
|
|
deprecation_version="1.12", |
|
|
removal_version="1.13", |
|
|
old_functional_name="routed_decode", |
|
|
) |
|
|
|
|
|
def add_handler(self, *handler: Callable) -> None: |
|
|
self.decoder.add_handler(*handler) |
|
|
|
|
|
def __iter__(self) -> Iterator[Tuple[str, Any]]: |
|
|
for data in self.datapipe: |
|
|
pathname = data[0] |
|
|
result = self.decoder(data) |
|
|
yield (pathname, result[pathname]) |
|
|
|
|
|
def __len__(self) -> int: |
|
|
if isinstance(self.datapipe, Sized): |
|
|
return len(self.datapipe) |
|
|
raise TypeError("{} instance doesn't have valid length".format(type(self).__name__)) |
|
|
|
