Update app.py

app.py

@@ -1,12 +1,11 @@
 from detect_face import detect_face
 from transformers import AutoModelForImageClassification
-from transformers import AutoImageProcessor
+from transformers import AutoFeatureExtractor
 from PIL import Image
 import tempfile
 import torch
 import gradio as gr
 from extract_frames import extract_frames
-from predict_image import predict_image
 
 import os
 import shutil
@@ -16,10 +15,37 @@ model = AutoModelForImageClassification.from_pretrained(
     "Jabrave/deepfake-detector"
 )
 
-processor = AutoImageProcessor.from_pretrained(
+processor = AutoFeatureExtractor.from_pretrained(
     "Jabrave/deepfake-detector"
 )
 
+face_model = AutoModelForImageClassification.from_pretrained(
+    "Jabrave/face-detector"
+)
+
+face_processor = AutoFeatureExtractor.from_pretrained(
+    "Jabrave/face-detector"
+)
+
+def predict_with_model(image, model, processor):
+
+    inputs = processor(images=image, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1).item()
+
+    confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()
+
+    label = model.config.id2label[predicted_class]
+
+    return {
+        "label": label,
+        "confidence": round(confidence * 100, 2)
+    }
+
 # โหลด labels จาก config อัตโนมัติ
 id2label = model.config.id2label
 
@@ -37,21 +63,39 @@ def predict(image):
     # ----------------------
     # วิเคราะห์ทั้งภาพ
     # ----------------------
-    full_result = predict_image(temp_path)
+    full_image = Image.open(temp_path)
+
+    full_result = predict_with_model(
+        full_image,
+        model,
+        processor
+    )
 
     # ----------------------
     # detect face
     # ----------------------
     os.makedirs("faces", exist_ok=True)
 
-    faces = detect_face(temp_path)
+    face_image = Image.open(face_path)
+
+    face_result = predict_with_model(
+        face_image,
+        face_model,
+        face_processor
+)
 
     face_scores = []
     fake_face_found = False
 
     for face_path in faces:
 
-        face_result = predict_image(face_path)
+        face_image = Image.open(face_path)
+
+        face_result = predict_with_model(
+            face_image,
+            face_model,
+            face_processor
+)
 
         face_scores.append(face_result["confidence"])
 
@@ -126,14 +170,25 @@ def predict_video(video_path):
         # ----------------------
         # วิเคราะห์ทั้งภาพ
         # ----------------------
-        full_result = predict_image(frame_path)
+        frame_image = Image.open(frame_path)
 
+        full_result = predict_with_model(
+            frame_image,
+            model,
+            processor
+)
         full_scores.append(full_result["confidence"])
 
         # ----------------------
         # detect หน้า
         # ----------------------
-        faces = detect_face(frame_path)
+        face_image = Image.open(face_path)
+
+        face_result = predict_with_model(
+            face_image,
+            face_model,
+            face_processor
+)
 
         face_fake_found = False