Update app.py

app.py

@@ -2,6 +2,7 @@ from detect_face import detect_face
 from transformers import AutoModelForImageClassification
 from transformers import AutoImageProcessor
 from PIL import Image
+import tempfile
 import torch
 import gradio as gr
 from extract_frames import extract_frames
@@ -26,31 +27,73 @@ print("Loaded labels:", id2label)
 
 def predict(image):
 
-    # แปลงเป็น PIL Image
-    image = Image.fromarray(image)
+    # ----------------------
+    # save temp image
+    # ----------------------
+    temp_path = "temp_image.jpg"
+
+    Image.fromarray(image).save(temp_path)
+
+    # ----------------------
+    # วิเคราะห์ทั้งภาพ
+    # ----------------------
+    full_result = predict_image(temp_path)
+
+    # ----------------------
+    # detect face
+    # ----------------------
+    os.makedirs("faces", exist_ok=True)
+
+    faces = detect_face(temp_path)
+
+    face_scores = []
+    fake_face_found = False
+
+    for face_path in faces:
+
+        face_result = predict_image(face_path)
+
+        face_scores.append(face_result["confidence"])
 
-    # preprocess
-    inputs = processor(images=image, return_tensors="pt")
+        if face_result["label"] != "real":
+            fake_face_found = True
 
-    # inference
-    with torch.no_grad():
-        outputs = model(**inputs)
+    # ----------------------
+    # combine score
+    # ----------------------
+    full_score = full_result["confidence"]
+
+    avg_face_score = (
+        sum(face_scores) / len(face_scores)
+        if face_scores else full_score
+    )
 
-    # softmax
-    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    final_score = (
+        full_score + avg_face_score
+    ) / 2
 
-    # class ที่มั่นใจสุด
-    pred = probs.argmax().item()
+    final_label = (
+        "artificial"
+        if (
+            full_result["label"] != "real"
+            or fake_face_found
+        )
+        else "real"
+    )
 
-    # confidence
-    confidence = probs[0][pred].item()
+    # cleanup
+    if os.path.exists(temp_path):
+        os.remove(temp_path)
 
-    # label จริงจาก model
-    label = id2label[pred]
+    if os.path.exists("faces"):
+        shutil.rmtree("faces")
 
     return {
-        "label": label,
-        "confidence": round(confidence * 100, 2)
+        "label": final_label,
+        "final_score": round(final_score, 2),
+        "full_image_score": round(full_score, 2),
+        "face_score": round(avg_face_score, 2),
+        "faces_detected": len(faces)
     }
 
 def predict_video(video_path):