Update app.py

app.py

@@ -1,26 +1,56 @@
-import gradio as gr
-from transformers import ViTForImageClassification, ViTImageProcessor
+from flask import Flask, request, jsonify
 from PIL import Image
+import requests
+from io import BytesIO
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# Load a simple AI model
-model_name = "google/vit-base-patch16-224"
-processor = ViTImageProcessor.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
+app = Flask(__name__)
 
-def detect_deepfake(image):
-    # Check if image is fake
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    return "REAL" if outputs.logits.argmax() == 0 else "FAKE (Deepfake!)"
+# Load model and processor (cache them for better performance)
+MODEL_NAME = "dima806/deepfake_vs_real_image_detection"
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
 
-# Create a simple web interface
-app = gr.Interface(
-    fn=detect_deepfake,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Label(label="Result"),
-    title="🔍 Deepfake Detector",
-    description="Upload a face photo to check if it's real or AI-generated!"
-)
+@app.route('/detect', methods=['POST'])
+def detect_deepfake():
+    try:
+        # Get image from request
+        if 'file' not in request.files:
+            return jsonify({"error": "No file uploaded"}), 400
+            
+        file = request.files['file']
+        if file.filename == '':
+            return jsonify({"error": "No selected file"}), 400
+            
+        # Open and process image
+        image = Image.open(file.stream).convert("RGB")
+        
+        # Preprocess image
+        inputs = processor(images=image, return_tensors="pt")
+        
+        # Run inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            
+        # Get probabilities
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        real_prob = round(probabilities[0][0].item() * 100, 2)
+        fake_prob = round(probabilities[0][1].item() * 100, 2)
+        
+        # Determine result
+        result = "Real" if real_prob > fake_prob else "Fake"
+        confidence = real_prob if result == "Real" else fake_prob
+        
+        return jsonify({
+            "result": result,
+            "confidence": confidence,
+            "real_probability": real_prob,
+            "fake_probability": fake_prob
+        })
+        
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
-# Launch the app with a public URL
-app.launch(share=True)  # << This gives you a temporary public link!
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)