Update demo.py

demo.py

@@ -5,114 +5,124 @@ from gradio.components import Image
 from PIL import Image as PILImage, ImageDraw, ImageFont  # This import may be needed if you're processing images
 from PIL import Image
 
+def face_crop(image, face_rect):
+    x1 = face_rect.get('x1')
+    y1 = face_rect.get('y1')
+    x2 = face_rect.get('x2')
+    y2 = face_rect.get('y2')
+    width = x2 - x1 + 1
+    height = y2 - y1 + 1
+
+
+    if x1 < 0:
+        x1 = 0
+    if y1 < 0:
+        y1 = 0
+    if x2 >= image.width:
+        x2 = image.width - 1
+    if y2 >= image.height:
+        y2 = image.height - 1
+
+    face_image = image.crop((x1, y1, x2, y2))
+    face_image_ratio = face_image.width / float(face_image.height)
+    resized_w = int(face_image_ratio * 150)
+    resized_h = 150
+
+    face_image = face_image.resize((int(resized_w), int(resized_h)))
+    return face_image
+
+def pil_image_to_base64(image, format="PNG"):
+    """
+    Converts a PIL.Image object to a Base64-encoded string.
+    :param image: PIL.Image object
+    :param format: Format to save the image, e.g., "PNG", "JPEG"
+    :return: Base64-encoded string
+    """
+    # Save the image to a BytesIO buffer
+    buffer = io.BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)  # Rewind the buffer
+
+    # Convert the buffer's contents to a Base64 string
+    base64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return base64_string
+    
 def compare_face(frame1, frame2):
     url = "http://127.0.0.1:8080/compare_face"
     files = {'file1': open(frame1, 'rb'), 'file2': open(frame2, 'rb')}
 
-    r = requests.post(url=url, files=files)
-
-    html = None
-    faces = None
-
-    compare_result = r.json().get('compare_result')
-    compare_similarity = r.json().get('compare_similarity')
-
-    html = ("<table>"
-                "<tr>"
-                    "<th>Compare Result</th>"
-                    "<th>Value</th>"
-                "</tr>"
-                "<tr>"
-                    "<td>Result</td>"
-                    "<td>{compare_result}</td>"
-                "</tr>"
-                "<tr>"
-                    "<td>Similarity</td>"
-                    "<td>{compare_similarity}</td>"
-                "</tr>"
-                "</table>".format(compare_result=compare_result, compare_similarity=compare_similarity))
-
+    file1 = None
+    file2 = None
     try:
-        image1 = Image.open(frame1)
-        image2 = Image.open(frame2)
-
-        face1 = None
-        face2 = None
-
-        if r.json().get('face1') is not None:
-            face = r.json().get('face1')
-            x1 = face.get('x1')
-            y1 = face.get('y1')
-            x2 = face.get('x2')
-            y2 = face.get('y2')
-
-            if x1 < 0:
-                x1 = 0
-            if y1 < 0:
-                y1 = 0
-            if x2 >= image1.width:
-                x2 = image1.width - 1
-            if y2 >= image1.height:
-                y2 = image1.height - 1
-
-            face1 = image1.crop((x1, y1, x2, y2))
-            face_image_ratio = face1.width / float(face1.height)
-            resized_w = int(face_image_ratio * 150)
-            resized_h = 150
-
-            face1 = face1.resize((int(resized_w), int(resized_h)))
-
-        if r.json().get('face2') is not None:
-            face = r.json().get('face2')
-            x1 = face.get('x1')
-            y1 = face.get('y1')
-            x2 = face.get('x2')
-            y2 = face.get('y2')
-
-            if x1 < 0:
-                x1 = 0
-            if y1 < 0:
-                y1 = 0
-            if x2 >= image2.width:
-                x2 = image2.width - 1
-            if y2 >= image2.height:
-                y2 = image2.height - 1
-
-            face2 = image2.crop((x1, y1, x2, y2))
-            face_image_ratio = face2.width / float(face2.height)
-            resized_w = int(face_image_ratio * 150)
-            resized_h = 150
-
-            face2 = face2.resize((int(resized_w), int(resized_h)))
-
-        if face1 is not None and face2 is not None:
-            new_image = Image.new('RGB',(face1.width + face2.width + 10, 150), (80,80,80))
-
-            new_image.paste(face1,(0,0))
-            new_image.paste(face2,(face1.width + 10, 0))
-            faces = new_image.copy()
-        elif face1 is not None and face2 is None:
-            new_image = Image.new('RGB',(face1.width + face1.width + 10, 150), (80,80,80))
-
-            new_image.paste(face1,(0,0))
-            faces = new_image.copy()
-        elif face1 is None and face2 is not None:
-            new_image = Image.new('RGB',(face2.width + face2.width + 10, 150), (80,80,80))
-
-            new_image.paste(face2,(face2.width + 10, 0))
-            faces = new_image.copy()
-
+        file1 = open(frame1, 'rb')
     except:
-        pass
+        return "Failed to open image1"
 
-    return [faces, html]
+    try:
+        file2 = open(frame2, 'rb')
+    except:
+        return "Failed to open image2"
+    
+    url = "http://127.0.0.1:8080/compare_face"
+    files = {'file1': file1, 'file2': file2}
+    result = requests.post(url=url, files=files)
+
+    if result.ok:
+        json_result = result.json()
+       
+        if json_result.get("resultCode") != "Ok":
+            return json_result.get("resultCode")
+        
+        try:
+            image1 = Image.open(frame1)
+            image2 = Image.open(frame2)
+
+            html = ""
+            faces1 = json_result.get("faces1", {})           
+            faces2 = json_result.get("faces2", {})
+            results = json_result.get("results", {})
+    
+            for result in results:
+                similarity = result.get('similarity')
+                face1_idx = result.get('face1')
+                face2_idx = result.get('face2')
+    
+                face_image1 = face_crop(image1, faces1[face1_idx])
+                face_value1 = ('<img src="data:image/png;base64,{base64_image}" style="width: 100px; height: auto; object-fit: contain;"/>').format(base64_image=pil_image_to_base64(face_image1, format="PNG"))
+    
+                face_image2 = face_crop(image2, faces2[face2_idx])
+                face_value2 = ('<img src="data:image/png;base64,{base64_image}" style="width: 100px; height: auto; object-fit: contain;"/>').format(base64_image=pil_image_to_base64(face_image2, format="PNG"))
+    
+                match_icon = '<svg fill="red" width="19" height="32" viewBox="0 0 19 32"><path d="M0 13.92V10.2H19V13.92H0ZM0 21.64V17.92H19V21.64H0Z"></path><path d="M14.08 0H18.08L5.08 32H1.08L14.08 0Z"></path></svg>'
+                if similarity > 0.67:
+                    match_icon = '<svg fill="green" width="19" height="32" viewBox="0 0 19 32"><path d="M0 13.9202V10.2002H19V13.9202H0ZM0 21.6402V17.9202H19V21.6402H0Z"></path></svg>'
+    
+                item_value = ('<div style="align-items: center; gap: 10px; display: flex; flex-direction: column;">'
+                                '<div style="display: flex; align-items: center; gap: 20px;">'
+                                '{face_value1}'
+                                '{match_icon}'
+                                '{face_value2}'
+                                '</div>'
+                                '<div style="text-align: center; margin-top: 10px;">'
+                                'Similarity: {similarity}'
+                                '</div>'
+                                '</div>'
+                ).format(face_value1=face_value1, face_value2=face_value2, match_icon=match_icon, similarity=f"{similarity:.2f}")
+
+                html += item_value
+                html += '<hr style="border: 1px solid #C0C0C0; margin: 10px 0;"/>'
+    
+            return html
+        except:
+            return "Processing failed"
+    else:
+        return result.text
 
 with gr.Blocks() as demo:
     gr.Markdown(
         """
     # KBY-AI - Face Recognition
     We offer SDKs for face recognition, liveness detection(anti-spoofing) and ID card recognition.
-
     <h4 style="display: flex; align-items: center;">
         ID Document Liveness Detection - Linux - <a href="https://web.kby-ai.com">https://web.kby-ai.com</a>
         <span>
@@ -135,21 +145,21 @@ with gr.Blocks() as demo:
         ```
         """
         )
-        with gr.Row():
-            with gr.Column():
-                compare_face_input1 = gr.Image(type='filepath')
-                gr.Examples(['face_examples/1.jpg', 'face_examples/3.jpg', 'face_examples/5.jpg', 'face_examples/7.jpg', 'face_examples/9.jpg'], 
-                            inputs=compare_face_input1)
-                compare_face_button = gr.Button("Compare Face")
-            with gr.Column():
-                compare_face_input2 = gr.Image(type='filepath')
-                gr.Examples(['face_examples/2.jpg', 'face_examples/4.jpg', 'face_examples/6.jpg', 'face_examples/8.jpg', 'face_examples/10.jpg'], 
-                            inputs=compare_face_input2)
-            with gr.Column():
-                compare_face_output = gr.Image(type="pil").style(height=150)
-                compare_result_output = gr.HTML(label='Result')
-
-        compare_face_button.click(compare_face, inputs=[compare_face_input1, compare_face_input2], outputs=[compare_face_output, compare_result_output])
+            with gr.Column(scale=7):
+                with gr.Row():
+                    with gr.Column():
+                        image_input1 = gr.Image(type='filepath')
+                        gr.Examples(['face_examples/1.jpg', 'face_examples/3.jpg', 'face_examples/5.jpg', 'face_examples/7.jpg', 'face_examples/9.jpg'], 
+                            inputs=image_input1)
+                    with gr.Column():
+                        image_input2 = gr.Image(type='filepath')
+                        gr.Examples(['face_examples/2.jpg', 'face_examples/4.jpg', 'face_examples/6.jpg', 'face_examples/8.jpg', 'face_examples/10.jpg'], 
+                            inputs=image_input2)
+                face_recog_button = gr.Button("Compare Face", variant="primary", size="lg")
+            with gr.Column(scale=3):
+                recog_html_output = gr.HTML()
+
+        face_recog_button.click(compare_face, inputs=[image_input1, image_input2], outputs=recog_html_output)
     gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fkby-ai%2FFaceRecognition"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fkby-ai%2FFaceRecognition&countColor=%23263759" /></a>')
 
 demo.launch(server_name="0.0.0.0", server_port=7860, share=True)