You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
33 lines
1.2 KiB
33 lines
1.2 KiB
import os
|
|
|
|
from .parser_image_folder import ParserImageFolder
|
|
from .parser_image_in_tar import ParserImageInTar
|
|
|
|
|
|
def create_parser(name, root, split='train', **kwargs):
|
|
name = name.lower()
|
|
name = name.split('/', 2)
|
|
prefix = ''
|
|
if len(name) > 1:
|
|
prefix = name[0]
|
|
name = "/".join(name[1:])
|
|
|
|
# FIXME improve the selection right now just tfds prefix or fallback path, will need options to
|
|
# explicitly select other options shortly
|
|
if prefix == 'tfds':
|
|
from .parser_tfds import ParserTfds # defer tensorflow import
|
|
parser = ParserTfds(root, name, split=split, **kwargs)
|
|
elif prefix == 'wds':
|
|
from .parser_wds import ParserWebdataset
|
|
kwargs.pop('download', False)
|
|
parser = ParserWebdataset(root, name, split=split, **kwargs)
|
|
else:
|
|
assert os.path.exists(root)
|
|
# default fallback path (backwards compat), use image tar if root is a .tar file, otherwise image folder
|
|
# FIXME support split here, in parser?
|
|
if os.path.isfile(root) and os.path.splitext(root)[1] == '.tar':
|
|
parser = ParserImageInTar(root, **kwargs)
|
|
else:
|
|
parser = ParserImageFolder(root, **kwargs)
|
|
return parser
|