Spaces:
Runtime error
Runtime error
File size: 1,325 Bytes
ec0fdfd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import os
import os.path
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset(path_files):
if path_files.find('.txt') != -1:
paths, size = make_dataset_txt(path_files)
else:
paths, size = make_dataset_dir(path_files)
return paths, size
def make_dataset_txt(files):
"""
:param path_files: the path of txt file that store the image paths
:return: image paths and sizes
"""
img_paths = []
with open(files) as f:
paths = f.readlines()
for path in paths:
path = path.strip()
if is_image_file(path) and os.path.exists(path):
img_paths.append(path)
return img_paths, len(img_paths)
def make_dataset_dir(dir):
"""
:param dir: directory paths that store the image
:return: image paths and sizes
"""
img_paths = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for root, _, fnames in os.walk(dir):
for fname in sorted(fnames):
if is_image_file(fname):
path = os.path.join(root, fname)
img_paths.append(path)
return img_paths, len(img_paths)
|