Update app.py

app.py

@@ -1,10 +1,12 @@
 import os
 import pathlib
+import glob
 
 import random
 from PIL import Image
+import plotly.graph_objects as go
+
 import torch
-import glob
 import numpy as np
 from basicsr.utils import imwrite
 from gfpgan import GFPGANer
@@ -31,27 +33,28 @@ def enhancer(image):
     cropped_faces, restored_faces, restored_img = restorer.enhance(
         image, has_aligned=False, only_center_face=False, paste_back=True)
 
-    return Image.fromarray(restored_faces[0]), np.array(restored_img)
+    return Image.fromarray(restored_faces[0])
 
-def inference(img1, img2):
+def inference(img1, img2, tolerance):
 
     img1 = face_recognition.face_encodings(img1)[0]
     img2 = face_recognition.face_encodings(img2)[0]
 
-    result = face_recognition.compare_faces([img1], img2)
-
-    if result[0]:
+    confidence = face_recognition.face_distance([img1], img2)[0]
+    
+    if confidence <= tolerance:
         result = 'People in the two images are same!'
     else:
         result = 'People in the two images are different!'
-    
-    return result
+        
+    return 1-confidence, result
 
 def main():
     
     st.set_page_config(page_title='QFace.ai', page_icon='🤖', layout='centered')
-    st.title('QFace.ai')
+    st.title('Qface.ai')
     col1, col2 = st.columns(2)
+    tolerance = 0.6
     
     with st.form('Input Form'):
         
@@ -70,8 +73,8 @@ def main():
                 
                 with col1:
                     img1 = Image.open(img1)
-                    eimg1_face, eimg1 = enhancer(img1)
-                    image_comparison(img1, eimg1_face,
+                    eimg1 = enhancer(img1)
+                    image_comparison(img1, eimg1,
                                      width=350,
                                      label1='Before',
                                      label2='After',
@@ -79,16 +82,25 @@ def main():
                                      make_responsive=False)
                 with col2:
                     img2 = Image.open(img2)
-                    eimg2_face, eimg2 = enhancer(img2)
-                    image_comparison(img2, eimg2_face,
+                    eimg2 = enhancer(img2)
+                    image_comparison(img2, eimg2,
                                      width=350,
                                      label1='Before',
                                      label2='After',
                                      show_labels=True,
                                      make_responsive=False)
                     
-            result = inference(np.array(img1), np.array(img2))
-            st.write(result)
+            confidence, result = inference(np.array(eimg1), np.array(eimg2), tolerance)
+            st.subheader(result)
+            
+            fig = go.Figure(go.Indicator(domain = {'x': [0, 1], 'y': [0, 1]},
+                                         value = round(confidence, 3),
+                                         mode = 'gauge+number',
+                                         title = {'text': 'Confidence'},
+                                         delta = {'reference': 0.5},
+                                         gauge = {'axis': {'range': [None, 1]}}))
+            
+            st.plotly_chart(fig, use_container_width=True)
     
 
 if __name__ == '__main__':