Add image check and confidence calculation in

app.py

@@ -16,6 +16,9 @@ def get_sample_images():
 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)
@@ -25,7 +28,9 @@ def inference(image):
     # decode
     dict = load_dict(os.path.join(os.path.dirname(__file__), 'label.names'))
     res = decode_label(pred, dict)
-    return res
+    # get confidence
+    confidence = get_confidence(pred)
+    return res, confidence
 
 
 class TFliteDemo:
@@ -79,6 +84,21 @@ 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)
+
+
 def decode_label(mat, chars) -> str:
     # mat is the output of model
     # get char indices along best path
@@ -90,20 +110,25 @@ def decode_label(mat, chars) -> str:
     res = res.replace(' ', '').replace('_', '')
     return res
 
-_TITLE = '''South Korean License Plate Recognition'''
-_DESCRIPTION = '''
-<div>
-<p style="text-align: center; font-size: 1.3em">This is a demo of South Korean License Plate Recognition.
-<a style="display:inline-block; margin-left: .5em" href='https://github.com/noahzhy/KR_LPR_TF/'><img src='https://img.shields.io/github/stars/noahzhy/KR_LPR_TF?style=social' /></a>
-</p>
-</div>
-'''
-interface = gr.Interface(
-    fn=inference,
-    inputs="image",
-    outputs="text",
-    title=_TITLE,
-    description=_DESCRIPTION,
-    examples=get_sample_images(),
-)
-interface.launch()
+
+if __name__ == '__main__':
+    _TITLE = '''South Korean License Plate Recognition'''
+    _DESCRIPTION = '''
+    <div>
+    <p style="text-align: center; font-size: 1.3em">This is a demo of South Korean License Plate Recognition.
+    <a style="display:inline-block; margin-left: .5em" href='https://github.com/noahzhy/KR_LPR_TF/'><img src='https://img.shields.io/github/stars/noahzhy/KR_LPR_TF?style=social' /></a>
+    </p>
+    </div>
+    '''
+    interface = gr.Interface(
+        fn=inference,
+        inputs="image",
+        outputs=[
+            gr.Textbox(label="Plate Number", type="text"),
+            gr.Textbox(label="Confidence", type="text"),
+        ],
+        title=_TITLE,
+        description=_DESCRIPTION,
+        examples=get_sample_images(),
+    )
+    interface.launch()