Spaces:
Sleeping
Sleeping
File size: 5,389 Bytes
e2ab2df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
# === 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))
|