folder.py

Raw

import os

import os . path

from pathlib import Path

from typing import Any , Callable , cast , Dict , List , Optional , Tuple , Union

from PIL import Image

from . vision import VisionDataset

def has_file_allowed_extension ( filename : str , extensions : Union [ str , Tuple [ str , ...]]) -> bool :

"""Checks if a file is an allowed extension.

Args:

filename (string): path to a file

extensions (tuple of strings): extensions to consider (lowercase)

Returns:

bool: True if the filename ends with one of given extensions

"""

return filename . lower (). endswith ( extensions if isinstance ( extensions , str ) else tuple ( extensions ))

def is_image_file ( filename : str ) -> bool :

"""Checks if a file is an allowed image extension.

Args:

filename (string): path to a file

Returns:

bool: True if the filename ends with a known image extension

"""

return has_file_allowed_extension ( filename , IMG_EXTENSIONS )

def find_classes ( directory : Union [ str , Path ]) -> Tuple [ List [ str ], Dict [ str , int ]]:

"""Finds the class folders in a dataset.

See :class:`DatasetFolder` for details.

"""

classes = sorted ( entry . name for entry in os . scandir ( directory ) if entry . is_dir ())

if not classes :

raise FileNotFoundError ( f"Couldn't find any class folder in { directory } ." )

class_to_idx = { cls_name : i for i , cls_name in enumerate ( classes )}

return classes , class_to_idx

def make_dataset (

directory : Union [ str , Path ],

class_to_idx : Optional [ Dict [ str , int ]] = None ,

extensions : Optional [ Union [ str , Tuple [ str , ...]]] = None ,

is_valid_file : Optional [ Callable [[ str ], bool ]] = None ,

allow_empty : bool = False ,

) -> List [ Tuple [ str , int ]]:

"""Generates a list of samples of a form (path_to_sample, class).

See :class:`DatasetFolder` for details.

Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function

by default.

"""

directory = os . path . expanduser ( directory )

if class_to_idx is None :

_ , class_to_idx = find_classes ( directory )

elif not class_to_idx :

raise ValueError ( "'class_to_index' must have at least one entry to collect any samples." )

both_none = extensions is None and is_valid_file is None

both_something = extensions is not None and is_valid_file is not None

if both_none or both_something :

raise ValueError ( "Both extensions and is_valid_file cannot be None or not None at the same time" )

if extensions is not None :

def is_valid_file ( x : str ) -> bool :

return has_file_allowed_extension ( x , extensions ) # type: ignore[arg-type]

is_valid_file = cast ( Callable [[ str ], bool ], is_valid_file )

instances = []

available_classes = set ()

for target_class in sorted ( class_to_idx . keys ()):

class_index = class_to_idx [ target_class ]

target_dir = os . path . join ( directory , target_class )

if not os . path . isdir ( target_dir ):

continue

for root , _ , fnames in sorted ( os . walk ( target_dir , followlinks = True )):

for fname in sorted ( fnames ):

path = os . path . join ( root , fname )

if is_valid_file ( path ):

item = path , class_index

instances . append ( item )

if target_class not in available_classes :

available_classes . add ( target_class )

empty_classes = set ( class_to_idx . keys ()) - available_classes

if empty_classes and not allow_empty :

msg = f"Found no valid file for the classes { ', ' . join ( sorted ( empty_classes )) } . "

if extensions is not None :

msg += f"Supported extensions are: { extensions if isinstance ( extensions , str ) else ', ' . join ( extensions ) } "

raise FileNotFoundError ( msg )

return instances

class DatasetFolder ( VisionDataset ):

"""A generic data loader.

This default directory structure can be customized by overriding the

:meth:`find_classes` method.

Args:

root (str or ``pathlib.Path``): Root directory path.

loader (callable): A function to load a sample given its path.

extensions (tuple[string]): A list of allowed extensions.

both extensions and is_valid_file should not be passed.

transform (callable, optional): A function/transform that takes in

a sample and returns a transformed version.

E.g, ``transforms.RandomCrop`` for images.

target_transform (callable, optional): A function/transform that takes

in the target and transforms it.

is_valid_file (callable, optional): A function that takes path of a file

and check if the file is a valid file (used to check of corrupt files)

both extensions and is_valid_file should not be passed.

allow_empty(bool, optional): If True, empty folders are considered to be valid classes.

An error is raised on empty folders if False (default).

Attributes:

classes (list): List of the class names sorted alphabetically.

class_to_idx (dict): Dict with items (class_name, class_index).

samples (list): List of (sample path, class_index) tuples

targets (list): The class_index value for each image in the dataset

"""

def __init__ (

self ,

root : Union [ str , Path ],

loader : Callable [[ str ], Any ],

extensions : Optional [ Tuple [ str , ...]] = None ,

transform : Optional [ Callable ] = None ,

target_transform : Optional [ Callable ] = None ,

is_valid_file : Optional [ Callable [[ str ], bool ]] = None ,

allow_empty : bool = False ,

) -> None :

super (). __init__ ( root , transform = transform , target_transform = target_transform )

classes , class_to_idx = self . find_classes ( self . root )

samples = self . make_dataset (

self . root ,

class_to_idx = class_to_idx ,

extensions = extensions ,

is_valid_file = is_valid_file ,

allow_empty = allow_empty ,

)

self . loader = loader

self . extensions = extensions

self . classes = classes

self . class_to_idx = class_to_idx

self . samples = samples

self . targets = [ s [ 1 ] for s in samples ]

@ staticmethod

def make_dataset (

directory : Union [ str , Path ],

class_to_idx : Dict [ str , int ],

extensions : Optional [ Tuple [ str , ...]] = None ,

is_valid_file : Optional [ Callable [[ str ], bool ]] = None ,

allow_empty : bool = False ,

) -> List [ Tuple [ str , int ]]:

"""Generates a list of samples of a form (path_to_sample, class).

This can be overridden to e.g. read files from a compressed zip file instead of from the disk.

Args:

directory (str): root dataset directory, corresponding to ``self.root``.

class_to_idx (Dict[str, int]): Dictionary mapping class name to class index.

extensions (optional): A list of allowed extensions.

Either extensions or is_valid_file should be passed. Defaults to None.

is_valid_file (optional): A function that takes path of a file

and checks if the file is a valid file

(used to check of corrupt files) both extensions and

is_valid_file should not be passed. Defaults to None.

allow_empty(bool, optional): If True, empty folders are considered to be valid classes.

An error is raised on empty folders if False (default).

Raises:

ValueError: In case ``class_to_idx`` is empty.

ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None.

FileNotFoundError: In case no valid file was found for any class.

Returns:

List[Tuple[str, int]]: samples of a form (path_to_sample, class)

"""

if class_to_idx is None :

# prevent potential bug since make_dataset() would use the class_to_idx logic of the

# find_classes() function, instead of using that of the find_classes() method, which

# is potentially overridden and thus could have a different logic.

raise ValueError ( "The class_to_idx parameter cannot be None." )

return make_dataset (

directory , class_to_idx , extensions = extensions , is_valid_file = is_valid_file , allow_empty = allow_empty

)

def find_classes ( self , directory : Union [ str , Path ]) -> Tuple [ List [ str ], Dict [ str , int ]]:

"""Find the class folders in a dataset structured as follows::

directory/

├── class_x

│ ├── xxx.ext

│ ├── xxy.ext

│ └── ...

│ └── xxz.ext

└── class_y

├── 123.ext

├── nsdf3.ext

└── ...

└── asd932_.ext

This method can be overridden to only consider

a subset of classes, or to adapt to a different dataset directory structure.

Args:

directory(str): Root directory path, corresponding to ``self.root``

Raises:

FileNotFoundError: If ``dir`` has no class folders.

Returns:

(Tuple[List[str], Dict[str, int]]): List of all classes and dictionary mapping each class to an index.

"""

return find_classes ( directory )

def __getitem__ ( self , index : int ) -> Tuple [ Any , Any ]:

"""

Args:

index (int): Index

Returns:

tuple: (sample, target) where target is class_index of the target class.

"""

path , target = self . samples [ index ]

sample = self . loader ( path )

if self . transform is not None :

sample = self . transform ( sample )

if self . target_transform is not None :

target = self . target_transform ( target )

return sample , target

def __len__ ( self ) -> int :

return len ( self . samples )

IMG_EXTENSIONS = ( ".jpg" , ".jpeg" , ".png" , ".ppm" , ".bmp" , ".pgm" , ".tif" , ".tiff" , ".webp" )

def pil_loader ( path : str ) -> Image . Image :

# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)

with open ( path , "rb" ) as f :

img = Image . open ( f )

return img . convert ( "RGB" )

# TODO: specify the return type

def accimage_loader ( path : str ) -> Any :

import accimage

try :

return accimage . Image ( path )

except OSError :

# Potentially a decoding problem, fall back to PIL.Image

return pil_loader ( path )

def default_loader ( path : str ) -> Any :

from torchvision import get_image_backend

if get_image_backend () == "accimage" :

return accimage_loader ( path )

else :

return pil_loader ( path )

class ImageFolder ( DatasetFolder ):

"""A generic data loader where the images are arranged in this way by default: ::

root/dog/xxx.png

root/dog/xxy.png

root/dog/[...]/xxz.png

root/cat/123.png

root/cat/nsdf3.png

root/cat/[...]/asd932_.png

This class inherits from :class:`~torchvision.datasets.DatasetFolder` so

the same methods can be overridden to customize the dataset.

Args:

root (str or ``pathlib.Path``): Root directory path.

transform (callable, optional): A function/transform that takes in a PIL image

and returns a transformed version. E.g, ``transforms.RandomCrop``

target_transform (callable, optional): A function/transform that takes in the

target and transforms it.

loader (callable, optional): A function to load an image given its path.

is_valid_file (callable, optional): A function that takes path of an Image file

and check if the file is a valid file (used to check of corrupt files)

allow_empty(bool, optional): If True, empty folders are considered to be valid classes.

An error is raised on empty folders if False (default).

Attributes:

classes (list): List of the class names sorted alphabetically.

class_to_idx (dict): Dict with items (class_name, class_index).

imgs (list): List of (image path, class_index) tuples

"""

def __init__ (

self ,

root : Union [ str , Path ],

transform : Optional [ Callable ] = None ,

target_transform : Optional [ Callable ] = None ,

loader : Callable [[ str ], Any ] = default_loader ,

is_valid_file : Optional [ Callable [[ str ], bool ]] = None ,

allow_empty : bool = False ,

super (). __init__ (

root ,

loader ,

IMG_EXTENSIONS if is_valid_file is None else None ,

transform = transform ,

target_transform = target_transform ,

is_valid_file = is_valid_file ,

allow_empty = allow_empty ,

)

self . imgs = self . samples

folder.py

Latest commit

History

folder.py

File metadata and controls