-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference.py
35 lines (28 loc) · 947 Bytes
/
inference.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
import cv2
import numpy as np
import tensorflow as tf
class PrintedWrittenClassifier():
def __init__(self, model_path):
self.model = tf.keras.models.load_model(model_path)
def preprocess(self, image):
IMG_SIZE = (256, 256)
im = cv2.resize(image, IMG_SIZE, interpolation = cv2.INTER_AREA)
im = np.expand_dims(im, axis = 0)
return im
def predict(self, img):
return self.model.predict(img)[0, 0]
def postprocess(self, y_pred):
# printed
if y_pred < 0.5:
return 0
# written
return 1
if __name__ == "__main__":
model_path = "./model.h5"
classifier = PrintedWrittenClassifier(model_path)
image_path = "test_images/printed.tif"
image = cv2.imread(image_path)
image = classifier.preprocess(image)
y_pred = classifier.predict(image)
output = classifier.postprocess(y_pred)
assert output == 0