-
Notifications
You must be signed in to change notification settings - Fork 0
/
cnn_blur_detection.py
93 lines (82 loc) · 3.24 KB
/
cnn_blur_detection.py
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
from keras.models import Sequential
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.optimizers import Adam
from sklearn.model_selection import train_test_split
from data_generator import get_regions_with_labels
from data_generator import BlurRegionsDataGenerator
from tensorflow import keras
from data_preparation import BLUR_COLS, BLUR_ROWS, get_img_regions
import numpy as np
class MotionBlurDetectionCNN:
def __init__(self, input_shape=(30, 30, 3), classes_number=2, epochs=10, learning_rate=0.01, batch_size=64,
test_size=0.2):
self.input_shape = input_shape
self.classes_number = classes_number
self.epochs = epochs
self.learning_rate = learning_rate
regions_train, regions_test = train_test_split(get_regions_with_labels(), test_size=test_size)
# Parameters for data generator
self.params = {
'dim': input_shape,
'batch_size': batch_size,
'n_classes': classes_number
}
# Generators
self.training_generator = BlurRegionsDataGenerator(regions_train, **self.params)
self.validation_generator = BlurRegionsDataGenerator(regions_test, **self.params)
self.model = Sequential(
[
Conv2D(96, (7, 7), input_shape=input_shape),
Activation('relu'),
MaxPooling2D(pool_size=(2, 2)),
Dropout(0.2),
Conv2D(256, (5, 5)),
Activation('relu'),
MaxPooling2D(pool_size=(2, 2)),
Dropout(0.2),
Flatten(),
Dense(1024),
Activation('relu'),
Dropout(0.2),
Dense(classes_number),
Activation('softmax')
]
)
Adam(lr=learning_rate)
self.model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=["accuracy"])
self.model.summary()
self.history = None
def get_model(self):
return self.model
def train_model(self):
hist = self.model.fit(
self.training_generator,
epochs=self.epochs,
validation_data=self.validation_generator,
use_multiprocessing=True,
workers=6,
verbose=1
)
self.history = hist
return hist
def predict(self, image_path):
mask = np.zeros((BLUR_ROWS, BLUR_COLS))
regions = get_img_regions(image_path)
crops = []
for region in regions:
crop = np.array([region.get_crop()])
prediciton = np.argmax(self.model.predict(crop, batch_size=1))
mask[region.row, region.col] = prediciton
return mask
def predict_crops(self, X_test):
probs = self.model.predict(X_test, batch_size=len(X_test))
return probs, np.argmax(probs, axis=1)
def test(self):
return self.model.evaluate(self.validation_generator)
def save_model(self, dir_path):
file = open(dir_path + "motionblur.h5", 'a')
self.model.save(dir_path + "motionblur.h5")
file.close()
def load_model(self, filename):
self.model = keras.models.load_model(filename)