Update app.py

app.py

@@ -4,30 +4,50 @@ import numpy as np
 from PIL import Image
 import json
 import uuid
+from pydrive2.auth import GoogleAuth
+from pydrive2.drive import GoogleDrive
 
-# --- Lưu ảnh tạm trong session (dictionary) ---
-image_storage = {}
+# --- Google Drive setup ---
+gauth = GoogleAuth()
+gauth.LocalWebserverAuth()  # mở trình duyệt để login
+drive = GoogleDrive(gauth)
 
-# 1️⃣ API: Upload ảnh
+# Dictionary lưu mapping image_id -> Google Drive file_id
+image_drive_map = {}
+
+# 1️⃣ Upload ảnh lên Google Drive
 def upload_image(image):
     try:
         if isinstance(image, Image.Image):
             image = np.array(image)
         image_id = str(uuid.uuid4())
-        image_storage[image_id] = image
+        filename = f"{image_id}.png"
+        pil_image = Image.fromarray(image)
+        pil_image.save(filename)  # tạm lưu local
+
+        gfile = drive.CreateFile({'title': filename})
+        gfile.SetContentFile(filename)
+        gfile.Upload()
+
+        image_drive_map[image_id] = gfile['id']
         return f"Ảnh lưu thành công với ID: {image_id}", image_id
     except Exception as e:
         return f"Lỗi khi upload: {str(e)}", None
 
-# 2️⃣ API: Nhận diện khuôn mặt
+# 2️⃣ Nhận diện khuôn mặt
 def recognize_face(image, image_id):
     try:
+        if image_id not in image_drive_map:
+            return "Ảnh ID không tồn tại"
         if isinstance(image, Image.Image):
             image = np.array(image)
-        if image_id not in image_storage:
-            return "Ảnh ID không tồn tại"
-        stored_image = image_storage[image_id]
-        result = DeepFace.verify(img1_path=image, img2_path=stored_image)
+
+        # Download ảnh từ Drive
+        file_id = image_drive_map[image_id]
+        gfile = drive.CreateFile({'id': file_id})
+        gfile.GetContentFile(f"temp_{image_id}.png")
+
+        result = DeepFace.verify(img1_path=image, img2_path=f"temp_{image_id}.png")
         output = {
             "Verified": result["verified"],
             "Khoảng cách": round(result["distance"], 4)
@@ -36,7 +56,7 @@ def recognize_face(image, image_id):
     except Exception as e:
         return f"Lỗi khi nhận diện: {str(e)}"
 
-# 3️⃣ API: Phân tích khuôn mặt
+# 3️⃣ Phân tích khuôn mặt
 def analyze_face(image):
     try:
         if isinstance(image, Image.Image):
@@ -57,14 +77,14 @@ def analyze_face(image):
     except Exception as e:
         return f"Lỗi khi phân tích: {str(e)}"
 
-# --- Tạo giao diện Gradio ---
+# --- Gradio interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("## 📁 Upload ảnh")
+    gr.Markdown("## 📁 Upload ảnh lên Google Drive")
     with gr.Row():
         upload_input = gr.Image(label="Upload ảnh", type="numpy")
         upload_btn = gr.Button("Upload")
         upload_output = gr.Textbox(label="Kết quả Upload")
-        image_id_store = gr.State()  # lưu image_id
+        image_id_store = gr.State()
 
     upload_btn.click(upload_image, inputs=upload_input, outputs=[upload_output, image_id_store])