recycling-image-classification/code/build_data_gens.py at master · bantucaravan/recycling-image-classification · GitHub

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
from tensorflow.keras.preprocessing.image import ImageDataGenerator


def get_gen_config(desc):
    '''
    Returns: Tuple of (train_gen_config, val_gen_config)
    '''
    return data_gen_configs[desc]


def list_variants():
    return list(data_gen_configs.keys())


data_gen_configs = {}


## !!! No name= attr in IMAGEDATAGENERATORS

######################################
######### No Augmentation

desc1 = 'No Augmentation'

train_gen_config1 = ImageDataGenerator(rescale=1./255)
val_gen_config1 = ImageDataGenerator(rescale=1./255)

data_gen_configs[desc1] = (train_gen_config1, val_gen_config1)


######################################
######### Baseline Augmentation

desc2 = 'Baseline Augmentation'

train_gen_config2 = ImageDataGenerator(rescale=1./255,
                                rotation_range=45,
                                width_shift_range=.15,
                                height_shift_range=.15,
                                #shear_range=.2,
                                horizontal_flip=True,
                                zoom_range=0.5,
                                fill_mode='nearest') #default
val_gen_config2 = ImageDataGenerator(rescale=1./255)


data_gen_configs[desc2] = (train_gen_config2, val_gen_config2)