deepfake_api.py

# === FASTAPI BACKEND (main.py) ===

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from PIL import Image
import io
import torch
import numpy as np
import cv2
import base64

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

device = 0 if torch.cuda.is_available() else -1

MODELS_CONFIG = {
    "SwinV2 Based": {"path": "haywoodsloan/ai-image-detector-deploy", "weight": 0.15},
    "ViT Based": {"path": "Heem2/AI-vs-Real-Image-Detection", "weight": 0.15},
    "SDXL Dataset": {"path": "Organika/sdxl-detector", "weight": 0.15},
    "SDXL + FLUX": {"path": "cmckinle/sdxl-flux-detector_v1.1", "weight": 0.15},
    "DeepFake v2": {"path": "prithivMLmods/Deep-Fake-Detector-v2-Model", "weight": 0.15},
    "Midjourney/SDXL": {"path": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL", "weight": 0.10},
    "ViT v4": {"path": "date3k2/vit-real-fake-classification-v4", "weight": 0.15},
}

models = {}
for name, config in MODELS_CONFIG.items():
    try:
        models[name] = pipeline("image-classification", model=config["path"], device=device)
    except Exception as e:
        print(f"Failed to load model {name}: {e}")

def pil_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode("utf-8")

def gen_ela(img_array, quality=90):
    if img_array.shape[2] == 4:
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    _, buffer = cv2.imencode('.jpg', img_array, encode_param)
    compressed_img = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
    ela_img = cv2.absdiff(img_array, compressed_img)
    ela_img = cv2.convertScaleAbs(ela_img, alpha=10)
    return Image.fromarray(cv2.cvtColor(ela_img, cv2.COLOR_BGR2RGB))

def gradient_processing(image_array):
    gray_img = cv2.cvtColor(image_array, cv2.COLOR_BGR2GRAY)
    dx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)
    dy = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=3)
    gradient_magnitude = cv2.magnitude(dx, dy)
    gradient_img = cv2.normalize(gradient_magnitude, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return Image.fromarray(gradient_img)

@app.post("/detect")
async def detect(image: UploadFile = File(...)):
    try:
        import time
        start_time = time.time()

        image_bytes = await image.read()
        input_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        img_np = np.array(input_image)
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

        individual_results = []
        weighted_ai_score = 0
        total_weight = 0

        aiModels = []
        colors = ["bg-red-500", "bg-orange-500", "bg-yellow-500", "bg-green-500", "bg-blue-500", "bg-purple-500", "bg-pink-500"]

        for i, (name, model_pipeline) in enumerate(models.items()):
            model_weight = MODELS_CONFIG[name]["weight"]
            predictions = model_pipeline(input_image)
            confidence = {p['label'].lower(): p['score'] for p in predictions}

            artificial_score = (
                confidence.get('artificial', 0) or confidence.get('ai image', 0) or
                confidence.get('ai', 0) or confidence.get('deepfake', 0) or
                confidence.get('ai_gen', 0) or confidence.get('fake', 0)
            )
            real_score = (
                confidence.get('real', 0) or confidence.get('real image', 0) or
                confidence.get('human', 0) or confidence.get('realism', 0)
            )

            if artificial_score > 0 and real_score == 0:
                real_score = 1.0 - artificial_score
            elif real_score > 0 and artificial_score == 0:
                artificial_score = 1.0 - real_score

            weighted_ai_score += artificial_score * model_weight
            total_weight += model_weight

            aiModels.append({
                "name": name,
                "percentage": round(artificial_score * 100, 2),
                "color": colors[i % len(colors)]
            })

        final_score = (weighted_ai_score / total_weight) * 100 if total_weight > 0 else 0
        verdict = final_score > 50
        processing_time = int((time.time() - start_time) * 1000)

        # Forensics
        ela_img = gen_ela(img_bgr)
        gradient_img = gradient_processing(img_bgr)

        return JSONResponse({
            "filename": image.filename,
            "isDeepfake": verdict,
            "confidence": round(final_score, 2),
            "aiModels": aiModels,
            "processingTime": processing_time,
            "forensics": {
                "original": pil_to_base64(input_image),
                "ela": pil_to_base64(ela_img),
                "gradient": pil_to_base64(gradient_img)
            },
            "verdictMessage": f"Consensus: {'Likely AI-Generated' if verdict else 'Likely Human-Made (Real)'}"
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



# === Add this new code to the bottom of your deepfake_api.py file ===
import os
import google.generativeai as genai
from pydantic import BaseModel, Field
from typing import List, Dict, Any

# Configure the Gemini API with the key from Hugging Face secrets
gemini_api_key = os.getenv("GOOGLE_API_KEY")
if gemini_api_key:
    genai.configure(api_key=gemini_api_key)

# Define Pydantic models to validate the incoming data structure
class AIModelResult(BaseModel):
    name: str
    percentage: float

class ForensicResult(BaseModel):
    original: str
    ela: str
    gradient: str

class DetectionResult(BaseModel):
    filename: str
    isDeepfake: bool
    confidence: float
    aiModels: List[AIModelResult]
    processingTime: int
    forensics: ForensicResult
    verdictMessage: str

# New endpoint to generate the report
@app.post("/generate-report")
async def generate_report(result: DetectionResult):
    if not gemini_api_key:
        raise HTTPException(status_code=500, detail="Google API key is not configured.")

    try:
        model = genai.GenerativeModel('gemini-1.5-flash')

        # Create a detailed prompt for the AI
        prompt = f"""
        You are an AI image forensics analyst. Your task is to generate a professional report based on the JSON data from a deepfake detection scan.

        The user has scanned the file: "{result.filename}".

        Here is the JSON data from the scan:
        {{
            "isDeepfake": {result.isDeepfake},
            "overall_confidence_score": {result.confidence}%,
            "verdict": "{result.verdictMessage}",
            "model_analysis": { {model.name: f"{model.percentage}%" for model in result.aiModels} }
        }}

        Please generate a report with the following structure using Markdown:

        ## Forensic Analysis Report: {result.filename}

        ### **Executive Summary**
        Provide a brief, high-level summary of the findings. State the final verdict clearly and the overall confidence score.

        ### **Detailed Model Analysis**
        Analyze the results from the individual AI models. Mention the models with the highest confidence scores (e.g., Midjourney/SDXL, SDXL Dataset) and explain what their findings imply.

        ### **Conclusion & Recommendation**
        Provide a final conclusion based on the combined evidence. Recommend next steps, such as whether the image can be trusted or if further manual analysis is required.
        """

        # Generate the report
        response = model.generate_content(prompt)

        return {"report": response.text}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to generate report: {str(e)}")