Niazi commited on
Commit
46cc8e7
·
1 Parent(s): a1b1ddb

Upload inference.py

Browse files
Files changed (1) hide show
  1. inference.py +29 -0
inference.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import keras
2
+
3
+ import numpy as np
4
+
5
+ class ModelInference:
6
+ """Inference module to predict an image class"""
7
+
8
+ def __init__(self, weights_path='./Weights/model.h5', threshold=0.8):
9
+ self.weights_path = weights_path
10
+ self.threshold = threshold
11
+
12
+ self.model = self.load_model()
13
+
14
+ def load_model(self):
15
+ return keras.models.load_model(self.weights_path)
16
+
17
+ def predict(self, image_array: np.array) -> bool:
18
+ model_output = self.model.predict(image_array)
19
+ prediction = self.parse_model_output(model_output)
20
+ return prediction
21
+
22
+ def parse_model_output(self, model_output: list) -> bool:
23
+ confidence_score = model_output[0][0]
24
+
25
+ result = False
26
+ if confidence_score >= self.threshold:
27
+ result = True
28
+
29
+ return result