Refactor inference function and add preprocessing

app.py

@@ -15,37 +15,20 @@ def get_sample_images():
     return [[i] for i in list_]
 
 
-def inference(image):
-    # load model
-    demo = TFliteDemo(os.path.join(os.path.dirname(__file__), 'model.tflite'))
-    # check image is not None
-    if image is None:
-        return 'None', 'None'
-    # load image
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    image = center_fit(image, 128, 64, top_left=True)
-    image = np.reshape(image, (1, *image.shape, 1)).astype(np.uint8)
-    # inference
-    pred = demo.inference(image)
-    # decode
-    dict = load_dict(os.path.join(os.path.dirname(__file__), 'label.names'))
-    res = decode_label(pred, dict)
-    # get confidence
-    confidence = get_confidence(pred)
-    return res, confidence
+def cv2_imread(path):
+    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
 
 
-class TFliteDemo:
-    def __init__(self, model_path):
-        self.interpreter = tf.lite.Interpreter(model_path=model_path)
-        self.interpreter.allocate_tensors()
-        self.input_details = self.interpreter.get_input_details()
-        self.output_details = self.interpreter.get_output_details()
-    
-    def inference(self, x):
-        self.interpreter.set_tensor(self.input_details[0]['index'], x)
-        self.interpreter.invoke()
-        return self.interpreter.get_tensor(self.output_details[0]['index'])
+def decode_label(mat, chars) -> str:
+    # mat is the output of model
+    # get char indices along best path
+    best_path_indices = np.argmax(mat[0], axis=-1)
+    # collapse best path (using itertools.groupby), map to chars, join char list to string
+    best_chars_collapsed = [chars[k] for k, _ in groupby(best_path_indices) if k != len(chars)]
+    res = ''.join(best_chars_collapsed)
+    # remove space and '_'
+    res = res.replace(' ', '').replace('_', '')
+    return res
 
 
 def center_fit(img, w, h, inter=cv2.INTER_NEAREST, top_left=True):
@@ -86,31 +69,60 @@ def load_dict(dict_path='label.names'):
     return dict
 
 
-def get_confidence(mat) -> float:
-    # mat is the output of model
-    # get char indices along best path
-    best_path_indices = np.argmax(mat[0], axis=-1)
-    confidence = np.max(mat[0], axis=-1)
-    blank_idx = mat.shape[-1] - 1
-    avg_confidence = []
-    for idx, conf in zip(best_path_indices, confidence):
-        if idx != blank_idx:
-            avg_confidence.append(conf)
-    conf = np.mean(avg_confidence) / 255.0
-    # keep 4 decimal places
-    return "{:.4f}".format(conf)
+class TFliteDemo:
+    def __init__(self, model_path, blank=85, conf_mode="mean"):
+        self.blank = blank
+        self.conf_mode = conf_mode
+        self.interpreter = tf.lite.Interpreter(model_path=model_path)
+        self.interpreter.allocate_tensors()
+        self.input_details = self.interpreter.get_input_details()
+        self.output_details = self.interpreter.get_output_details()
 
+    def inference(self, x):
+        self.interpreter.set_tensor(self.input_details[0]['index'], x)
+        self.interpreter.invoke()
+        return self.interpreter.get_tensor(self.output_details[0]['index'])
 
-def decode_label(mat, chars) -> str:
-    # mat is the output of model
-    # get char indices along best path
-    best_path_indices = np.argmax(mat[0], axis=-1)
-    # collapse best path (using itertools.groupby), map to chars, join char list to string
-    best_chars_collapsed = [chars[k] for k, _ in groupby(best_path_indices) if k != len(chars)]
-    res = ''.join(best_chars_collapsed)
-    # remove space and '_'
-    res = res.replace(' ', '').replace('_', '')
-    return res
+    def preprocess(self, img):
+        if isinstance(img, str):
+            image = cv2_imread(img)
+        else:
+            image = img
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        image = center_fit(image, 128, 64, top_left=True)
+        image = np.reshape(image, (1, *image.shape, 1)).astype(np.uint8)
+        return image
+
+    def get_confidence(self, pred, mode="mean"):
+        conf = []
+        idxs = np.argmax(pred, axis=-1)
+        values = np.max(pred, axis=-1)
+
+        for idx, c in zip(idxs, values):
+            if idx == self.blank: continue
+            conf.append(c/255)
+
+        if mode == "min":
+            return np.min(conf)
+
+        return np.mean(conf)
+
+    def postprocess(self, pred):
+        label = decode_label(pred, load_dict())
+        conf = self.get_confidence(pred[0], mode=self.conf_mode)
+        # keep 4 decimal places
+        conf = float('{:.4f}'.format(conf))
+        return label, conf
+
+
+def inference(img):
+    # preprocess
+    img = demo.preprocess(img)
+    # inference
+    pred = demo.inference(img)
+    # postprocess
+    label, conf = demo.postprocess(pred)
+    return label, conf
 
 
 if __name__ == '__main__':
@@ -122,6 +134,8 @@ if __name__ == '__main__':
     </p>
     </div>
     '''
+    # init model
+    demo = TFliteDemo(os.path.join(os.path.dirname(__file__), 'model.tflite'))
     interface = gr.Interface(
         fn=inference,
         inputs="image",