Update app.py

app.py

@@ -1,32 +1,33 @@
-from flask import Flask, render_template, request, jsonify
-from detect import DeepfakeDetector
-import os
+from transformers import ViTForImageClassification, ViTImageProcessor
+import torch
+from PIL import Image
+import gradio as gr
 
-app = Flask(__name__)
-detector = DeepfakeDetector()
+# Load pre-trained model and processor
+model_name = "facebook/deit-base-distilled-patch16-224"
+processor = ViTImageProcessor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
 
-@app.route("/")
-def home():
-    return render_template("index.html")
-
-@app.route("/detect", methods=["POST"])
-def detect():
-    if "file" not in request.files:
-        return jsonify({"error": "No file uploaded"}), 400
-    
-    file = request.files["file"]
-    if file.filename == "":
-        return jsonify({"error": "Empty filename"}), 400
+def detect_deepfake(image):
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
     
-    temp_path = f"temp_{file.filename}"
-    file.save(temp_path)
+    # Make prediction
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
     
-    try:
-        result = detector.detect_from_file(temp_path)
-        return jsonify({"result": result})
-    finally:
-        if os.path.exists(temp_path):
-            os.remove(temp_path)
+    # For demonstration, we'll assume class 0 is real and 1 is fake
+    # (In a real project, you'd need to verify this with your model)
+    return "Real" if predicted_class_idx == 0 else "Fake (Possible Deepfake)"
+
+# Create a simple interface
+iface = gr.Interface(
+    fn=detect_deepfake,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Deepfake Detection",
+    description="Upload an image to check if it might be a deepfake."
+)
 
-if __name__ == "__main__":
-    app.run(debug=True)
+iface.launch()