Andercar's picture
Update main.py
15c813b verified
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
from transformers import pipeline
app = FastAPI()
# Allow CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Load AI Detection Model
# We use a pre-trained model from Hugging Face
# 'Organika/sdxl-detector' is specialized for Stable Diffusion detection
print("Loading AI Model...")
classifier = pipeline("image-classification", model="Organika/sdxl-detector")
print("Model Loaded!")
@app.post("/analyze")
async def analyze_media(file: UploadFile = File(...)):
try:
contents = await file.read()
image = Image.open(io.BytesIO(contents))
# Run Inference
results = classifier(image)
# results is a list like [{'label': 'artificial', 'score': 0.99}, {'label': 'human', 'score': 0.01}]
# Find the 'artificial' or 'AI' score
ai_score = 0.0
for r in results:
label = r['label'].lower()
if 'artificial' in label or 'ai' in label:
ai_score = r['score']
break
if 'human' in label or 'real' in label:
# If we found human score, AI score is 1 - human
ai_score = 1.0 - r['score']
is_ai = ai_score > 0.5
return {
"filename": file.filename,
"is_ai": is_ai,
"confidence": round(ai_score * 100, 2), # Return 0-100
"details": results
}
except Exception as e:
print(f"Error: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
def read_root():
return {"status": "AI Detector Neural Network is Running"}