from flask import Flask, request, jsonify
from PIL import Image
import requests
from io import BytesIO
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

app = Flask(__name__)

# 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)

@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

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)