Update deepfake_api.py

deepfake_api.py

@@ -1,136 +1,213 @@
-# === 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))
+# === 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)}")