Update app.py

app.py

@@ -1,105 +1,124 @@
-# app.py
 import base64
 import cv2
 import numpy as np
 import requests
 from fastapi import FastAPI
 from pydantic import BaseModel
+import uvicorn
 import insightface
-import logging
+import gradio as gr
 
-# ---------------------------
-# Reduce InsightFace logging
-# ---------------------------
-logging.getLogger("insightface").setLevel(logging.ERROR)
-
-# ---------------------------
-# Load face detector + recognition model
-# ---------------------------
+# ---------- Load Face Detector + Recognition Model ----------
 model = insightface.app.FaceAnalysis(name="buffalo_l")
-model.prepare(ctx_id=-1, det_size=(640, 640))  # CPU only, avoid CUDA warnings
+model.prepare(ctx_id=0, det_size=(640, 640))
 
-# ---------------------------
-# FastAPI app
-# ---------------------------
-app = FastAPI(title="Face Compare API")
+# ---------- FastAPI App ----------
+app = FastAPI()
 
-# ---------------------------
-# Request schema
-# ---------------------------
+# ---------- API Request Schema ----------
 class CompareRequest(BaseModel):
-    image1: str   # base64 string OR URL
-    image2: str   # base64 string OR URL
-
-# ---------------------------
-# Helper: Load image from base64 or URL
-# ---------------------------
-def load_image(input_str):
-    if input_str.startswith("http://") or input_str.startswith("https://"):
-        try:
-            resp = requests.get(input_str, timeout=5)
-            if resp.status_code != 200:
-                return None
-            arr = np.frombuffer(resp.content, np.uint8)
-            img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
-            return img
-        except:
-            return None
-    else:
-        try:
-            if "," in input_str:
-                b64 = input_str.split(",", 1)[1]
-            else:
-                b64 = input_str
-            data = base64.b64decode(b64)
-            arr = np.frombuffer(data, np.uint8)
-            img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
-            return img
-        except:
-            return None
-
-# ---------------------------
-# Helper: Get face embedding
-# ---------------------------
-def get_embedding(img):
+    image1: str | None = None     # base64
+    image2: str | None = None     # base64
+    image1_url: str | None = None # URL
+    image2_url: str | None = None # URL
+
+# ---------- Helper: Convert base64 to CV2 image ----------
+def b64_to_img(b64_string):
+    try:
+        img_data = base64.b64decode(b64_string)
+        np_arr = np.frombuffer(img_data, np.uint8)
+        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        return img
+    except:
+        return None
+
+# ---------- Helper: Convert URL to CV2 image ----------
+def url_to_img(url):
     try:
-        faces = model.get(img)
-        if len(faces) == 0:
-            return None
-        return faces[0].embedding
+        resp = requests.get(url, timeout=5)
+        np_arr = np.frombuffer(resp.content, np.uint8)
+        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        return img
     except:
         return None
 
-# ---------------------------
-# Helper: Cosine similarity
-# ---------------------------
-def cosine_similarity(a, b):
-    a = np.array(a)
-    b = np.array(b)
-    denom = np.linalg.norm(a) * np.linalg.norm(b)
-    if denom == 0:
-        return 0.0
-    return float(np.dot(a, b) / denom)
-
-# ---------------------------
-# API endpoint
-# ---------------------------
+# ---------- Helper: Extract face embedding ----------
+def get_embedding(img):
+    faces = model.get(img)
+    if len(faces) == 0:
+        return None
+    return faces[0].embedding  # first detected face
+
+# ---------- POST /compare API ----------
 @app.post("/compare")
 async def compare_faces(req: CompareRequest):
-    img1 = load_image(req.image1)
-    img2 = load_image(req.image2)
+
+    # ---- Load Image 1 ----
+    if req.image1:
+        img1 = b64_to_img(req.image1)
+    elif req.image1_url:
+        img1 = url_to_img(req.image1_url)
+    else:
+        img1 = None
+
+    # ---- Load Image 2 ----
+    if req.image2:
+        img2 = b64_to_img(req.image2)
+    elif req.image2_url:
+        img2 = url_to_img(req.image2_url)
+    else:
+        img2 = None
 
     if img1 is None or img2 is None:
-        return {"error": "Cannot load one or both images"}
+        return {"error": "Invalid image data or URL."}
 
     emb1 = get_embedding(img1)
     emb2 = get_embedding(img2)
 
     if emb1 is None or emb2 is None:
-        return {"error": "No face detected in one or both images"}
+        return {"error": "No face detected in one or both images."}
+
+    # Cosine similarity
+    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
+
+    matched = similarity > 0.55  # threshold; adjust as needed
+
+    return {
+        "similarity": float(similarity),
+        "match": matched
+    }
+
+# ---------- Gradio UI ----------
+def gradio_ui(img1, img2):
+    import base64
+    import io
+
+    def to_b64(pil_img):
+        buf = io.BytesIO()
+        pil_img.save(buf, format="JPEG")
+        return base64.b64encode(buf.getvalue()).decode()
+
+    if img1 is None or img2 is None:
+        return "Upload both images."
+
+    # Convert to OpenCV
+    img1_cv = cv2.cvtColor(np.array(img1), cv2.COLOR_RGB2BGR)
+    img2_cv = cv2.cvtColor(np.array(img2), cv2.COLOR_RGB2BGR)
+
+    emb1 = get_embedding(img1_cv)
+    emb2 = get_embedding(img2_cv)
+
+    if emb1 is None or emb2 is None:
+        return "Face not detected."
+
+    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
+    matched = similarity > 0.55
 
-    similarity = cosine_similarity(emb1, emb2)
-    threshold = 0.55  # adjust threshold for stricter or looser matching
-    match = similarity >= threshold
+    return f"Similarity: {similarity:.3f} | Match: {matched}"
 
-    return {"similarity": similarity, "match": match}
+gr.Interface(
+    fn=gradio_ui,
+    inputs=[gr.Image(), gr.Image()],
+    outputs="text",
+    title="Face Match API"
+).launch()