ReflectionEraser/ReflectionEraserApp

1###############################################################################
2# Code from
3# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
4# Modified the original code so that it also loads images from the current
5# directory as well as the subdirectories
6###############################################################################
7
8import torch.utils.data as data
9
10from PIL import Image
11import os
12import os.path
13
14IMG_EXTENSIONS = [
15    '.jpg', '.JPG', '.jpeg', '.JPEG',
16    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
17]
18
19
20def read_fns(filename):#returns all lines in a file as a list
21    with open(filename) as f:
22        fns = f.readlines()
23        fns = [fn.strip() for fn in fns] #This strips any leading and trailing whitespace characters (including newlines) from each line in fns and turns result into a list
24    return fns
25
26
27def is_image_file(filename):#checks if file is image based on extensions
28    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
29
30
31def make_dataset(dir, fns=None):#makes dataset from directory with files
32    images = []
33    assert os.path.isdir(dir), '%s is not a valid directory' % dir
34
35    if fns is None:#if no filename list is given, create list of images
36        for root, _, fnames in sorted(os.walk(dir)): #walks through the directory tree starting from dir, returning the root directory path, directories, and file names. The results are sorted.
37            for fname in fnames:
38                if is_image_file(fname):                
39                    path = os.path.join(root, fname)
40                    images.append(path)
41    else:#create list of images from list
42        for fname in fns:
43            if is_image_file(fname):
44                path = os.path.join(dir, fname)
45                images.append(path)
46
47    return images
48
49
50def default_loader(path):#convert image in path to rgb and open it
51    return Image.open(path).convert('RGB') #if grayscale or has alpha channel (rgba)
52