CollabAR-Code/trainExpert.py at master · CollabAR-Source/CollabAR-Code · 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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import numpy as np
import cv2
import os
import sys
from keras.applications.mobilenetv2 import MobileNetV2
from keras.layers import Dense, Dropout, Activation, Flatten, Input
from keras.layers import Conv2D, MaxPooling2D
from keras.models import Model
from keras.models import Sequential
import keras
from keras.layers import Reshape
from skimage.util import random_noise
import random
from keras_preprocessing.image import ImageDataGenerator


def pristine(image):
    return image/255

def motion_blur(image):
    blurred = image
    if random.randint(1,1)== 1:
        degree = random.randint(0,40)
        if degree != 0:
            angle = random.randint(1,360)
            M = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
            motion_blur_kernel = np.diag(np.ones(degree))
            motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (degree, degree))

            motion_blur_kernel = motion_blur_kernel / degree
            blurred = cv2.filter2D(image, -1, motion_blur_kernel)
            cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
    if blurred.mean()>1:
        blurred = blurred/255
    return blurred

def gaussion_noise(image):
    if random.randint(1,1)== 1:
        k = random.randint(0,40)
        k = k /1000
        if image.mean()>1:
            image = image/255
        image = random_noise(image, mode="gaussian", var = k)
    if image.mean()>1:
        image = image/255
    return image

def gaussion_blur(image):
    if random.randint(1,1)==1:
        k = random.randint(0,20)
        if k!=0:
            k = 2*k+1
            image = cv2.GaussianBlur(image,(k,k),0)
    if image.mean()>1:
        image = image/255
    return image

def build_model( ):
    target_size = 224
    input_tensor = Input(shape=(target_size, target_size, 3))
    base_model = MobileNetV2(
        include_top=False,
        weights='imagenet',
        input_tensor=input_tensor,
        input_shape=(target_size, target_size, 3),
        pooling='avg')
    for layer in base_model.layers:
        layer.trainable = True  # trainable has to be false in order to freeze the layers
    op = Dense(256, activation='relu')(base_model.output)
    op = Dropout(.25)(op)
    output_tensor = Dense(257, activation='softmax')(op)
    model = Model(inputs=input_tensor, outputs=output_tensor)
    return model

def main(argv):
    folder = os.path.exists('./weights')
    if not folder:
        os.makedirs('./weights')
    batch_size = 32
    model = build_model()
    model.compile(optimizer=keras.optimizers.Adam(lr = 0.0001),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    if argv[1] == 'pristin':
        train_datagen = ImageDataGenerator(preprocessing_function=pristine)
        validation_datagen = ImageDataGenerator(preprocessing_function=pristine)
        test_datagen = ImageDataGenerator(preprocessing_function=pristine)
        filepath="./weights/pristine_expert.hdf5"
    elif argv[1] == 'MB':
        train_datagen = ImageDataGenerator(preprocessing_function=motion_blur)
        validation_datagen = ImageDataGenerator(preprocessing_function=motion_blur)
        test_datagen = ImageDataGenerator(preprocessing_function=motion_blur)
        filepath="./weights/motion_blur_expert.hdf5"
    elif argv[1] == 'GB':
        train_datagen = ImageDataGenerator(preprocessing_function=gaussion_blur)
        validation_datagen = ImageDataGenerator(preprocessing_function=gaussion_blur)
        test_datagen = ImageDataGenerator(preprocessing_function=gaussion_blur)
        filepath="./weights/Gaussian_blur_expert.hdf5"
    else:
        train_datagen = ImageDataGenerator(preprocessing_function=gaussion_noise)
        validation_datagen = ImageDataGenerator(preprocessing_function=gaussion_noise)
        test_datagen = ImageDataGenerator(preprocessing_function=gaussion_noise)
        filepath="./weights/Gaussian_noise_expert.hdf5"

    train_generator = train_datagen.flow_from_directory('train', target_size=(224,224),batch_size=batch_size, class_mode = 'categorical')
    validation_generator = validation_datagen.flow_from_directory('validation', target_size=(224,224),batch_size=batch_size, class_mode = 'categorical')
    test_generator = test_datagen.flow_from_directory('test', target_size=(224,224),batch_size=batch_size, class_mode = 'categorical')

    checkpoint = keras.callbacks.ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True, mode='auto')
    callbacks_list = [checkpoint]

    if argv[1] != 'pristin':
        model.load_weights('./weights/pristin_expert.hdf5')

    STEP_SIZE_TRAIN=train_generator.n//train_generator.batch_size
    STEP_SIZE_VALID=validation_generator.n//validation_generator.batch_size
    history = model.fit_generator(train_generator,
                        steps_per_epoch= STEP_SIZE_TRAIN ,
                        epochs=100,
                        validation_data=validation_generator,
                        validation_steps=STEP_SIZE_VALID,
                        callbacks = callbacks_list)

if __name__== "__main__":
    main(sys.argv)