Update demo.py

demo.py

@@ -3,123 +3,108 @@ import requests
 import datadog_api_client
 from PIL import Image
 
-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))
-
+def face_crop(image, face_rect):
+    x = face_rect.get('x')
+    y = face_rect.get('y')
+    width = face_rect.get('width')
+    height = face_rect.get('height')
+
+
+    if x < 0:
+        x = 0
+    if y < 0:
+        y = 0
+    if x + width >= image.width:
+        width = image.width - x
+    if y + height >= image.height:
+        height = image.height - y
+
+    face_image = image.crop((x, y, x + width - 1, y + height - 1))
+    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 compare_face(image1, image2):
     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()
-
+        img_bytes1 = io.BytesIO()
+        image1.save(img_bytes1, format="JPEG")
+        img_bytes1.seek(0)
     except:
-        pass
-
-    return [faces, html]
+        return ["Failed to open image1", {"resultCode": "Failed to open image1"}]
 
+    try:
+        img_bytes2 = io.BytesIO()
+        image2.save(img_bytes2, format="JPEG")
+        img_bytes2.seek(0)
+    except:
+        return ["Failed to open image2", {"resultCode": "Failed to open image2"}]
+    
+    url = "http://127.0.0.1:8080/compare_face"
+    files = {'image1': img_bytes1, 'image2': img_bytes2}
+    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"), json_result]
+
+        html = ""
+        faces1 = json_result.get("faces1", {})
+        faces2 = json_result.get("faces2", {})
+        results = json_result.get("results", {})
+
+        for result in results:
+            score = result.get('score')
+            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 score > 0.7:
+                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;">'
+                            'Score: {score}'
+                            '</div>'
+                            '</div>'
+            ).format(face_value1=face_value1, face_value2=face_value2, match_icon=match_icon, score=f"{score:.2f}")
+            html += item_value
+            html += '<hr style="border: 1px solid #C0C0C0; margin: 10px 0;"/>'
+
+        return html
+    else:
+        return result.text
+        
 with gr.Blocks(css=".gradio-container {background-color: #F4E5E0}") 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/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/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='pil')
+                    gr.Examples(['examples/1.webp', 'examples/2.webp', 'examples/3.webp', 'examples/4.webp'], 
+                                inputs=image_input1)
+                with gr.Column():
+                    image_input2 = gr.Image(type='pil')
+                    gr.Examples(['examples/5.webp', 'examples/6.webp', 'examples/7.webp', 'examples/8.webp'], 
+                                inputs=image_input2)
+            verifyThreshold = gr.Slider(minimum=0, maximum=1, value=0.67, label="Verify Threshold")    
+            face_recog_button = gr.Button("Face Recognition")
+        with gr.Column(scale=3):
+            recog_html_output = gr.HTML()
+
+    face_recog_button.click(compare_face, inputs=[image_input1, image_input2, verifyThreshold], outputs=recog_html_output)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)