KenzyPhill/DEEPFAKE_FORENSIC_LAB

1from fastapi import FastAPI, File, UploadFile
2from fastapi.middleware.cors import CORSMiddleware
3from fastapi.staticfiles import StaticFiles
4from fastapi.responses import FileResponse
5from transformers import pipeline
6from PIL import Image, ImageChops, ImageEnhance
7import io
8import base64
9import uvicorn
10import os
11
12app = FastAPI()
13
14# CORS settings 
15app.add_middleware(
16    CORSMiddleware,
17    allow_origins=["*"],
18    allow_credentials=True,
19    allow_methods=["*"],
20    allow_headers=["*"],
21)
22
23# Loading AI models
24print("--- Initializing AI Brains... ---")
25try:
26    model1 = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")
27    model2 = pipeline("image-classification", model="umm-maybe/AI-image-detector")
28    print("--- AI Brains Online and Ready! ---")
29except Exception as e:
30    print(f"Error loading models: {e}")
31
32# ELA Forensic Logic
33def get_ela_image(image_bytes, quality=90):
34    original = Image.open(io.BytesIO(image_bytes)).convert('RGB')
35    
36    # Save image with specific quality to check compression difference
37    temp_io = io.BytesIO()
38    original.save(temp_io, format='JPEG', quality=quality)
39    temp_io.seek(0)
40    resaved = Image.open(temp_io)
41    
42    # Calculate ELA
43    ela_image = ImageChops.difference(original, resaved)
44    
45    # Enhance the result so it's visible to human eye
46    extrema = ela_image.getextrema()
47    max_diff = max([ex[1] for ex in extrema])
48    if max_diff == 0: max_diff = 1
49    scale = 255.0 / max_diff
50    ela_image = ImageEnhance.Brightness(ela_image).enhance(scale)
51    
52    # Convert back to base64 for frontend
53    buffered = io.BytesIO()
54    ela_image.save(buffered, format="PNG")
55    return base64.b64encode(buffered.getvalue()).decode('utf-8')
56
57# Prediction Endpoint
58@app.post("/predict")
59async def predict_image(file: UploadFile = File(...)):
60    try:
61        contents = await file.read()
62        image = Image.open(io.BytesIO(contents)).convert("RGB")
63        
64        # Dual Model Analysis
65        res1 = model1(image)
66        res2 = model2(image)
67        
68        # Forensic Analysis
69        ela_base64 = get_ela_image(contents)
70        
71        return {
72            "status": "success", 
73            "model_1": res1,
74            "model_2": res2,
75            "ela_image": ela_base64
76        }
77    except Exception as e:
78        return {"status": "error", "message": str(e)}
79
80# Serve the HTML Frontend
81@app.get("/")
82async def read_index():
83    return FileResponse('index.html')
84
85# Start Server (Hugging Face default port is 7860)
86if __name__ == "__main__":
87    uvicorn.run(app, host="0.0.0.0", port=7860)