from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse import torch import numpy as np from PIL import Image from transformers import AutoImageProcessor, AutoModelForImageClassification import time import io import uvicorn # Initialize FastAPI app app = FastAPI(title="AI vs Real Image Detection API") # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Model configurations - Optimized weights based on model performance MODEL_CONFIGS = { "dima806/deepfake_vs_real_image_detection": { "name": "DeepFake Detector", "weight": 1.0, # Higher weight for primary deepfake detector }, "umm-maybe/AI-image-detector": { "name": "AI Image Detector", "weight": 0.8, # Adjusted weight for general AI detection }, "Organika/sdxl-detector": { "name": "SDXL Detector", "weight": 0.9, # Higher weight for SDXL-specific detection } } # Global variables for models models = {} processors = {} device = None def load_models(): """Load all models with enhanced error handling and memory optimization""" global models, processors, device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") success_count = 0 for model_id, config in MODEL_CONFIGS.items(): try: print(f"Loading {config['name']} ({model_id})...") # Load processor with error handling try: processor = AutoImageProcessor.from_pretrained( model_id, low_cpu_mem_usage=True ) print(f"✅ Processor loaded for {config['name']}") except Exception as proc_error: print(f"❌ Processor failed for {config['name']}: {proc_error}") continue # Load model with different strategies for CUDA vs CPU if device.type == "cuda": model = AutoModelForImageClassification.from_pretrained( model_id, torch_dtype=torch.float16, device_map="auto", trust_remote_code=False, # More secure ) model = model.to(device) model.half() else: # CPU-optimized loading for HuggingFace Spaces model = AutoModelForImageClassification.from_pretrained( model_id, torch_dtype=torch.float32, device_map={'': 'cpu'}, low_cpu_mem_usage=True, trust_remote_code=False, # More secure ) model.eval() processors[model_id] = processor models[model_id] = model success_count += 1 print(f"✅ {config['name']} loaded successfully") # Force garbage collection after each model import gc gc.collect() except Exception as e: print(f"❌ Failed to load {config['name']}: {str(e)}") print(f"Error type: {type(e).__name__}") continue print(f"✅ Model loading complete! Successfully loaded: {success_count}/{len(MODEL_CONFIGS)}") if success_count == 0: print("⚠️ WARNING: No models loaded successfully!") return False return True def preprocess_image(image, processor): """Enhanced preprocessing for better accuracy""" # Ensure RGB format if image.mode != 'RGB': image = image.convert('RGB') # Resize to optimal size if too large (reduces noise, improves accuracy) max_size = 512 if max(image.size) > max_size: ratio = max_size / max(image.size) new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio)) image = image.resize(new_size, Image.Resampling.LANCZOS) # Process with the model's processor inputs = processor(images=image, return_tensors="pt") return inputs def predict_single_model(image, model, processor, device): """Make prediction with a single model - enhanced accuracy""" try: inputs = preprocess_image(image, processor) # Move inputs to device if device.type == "cuda": inputs = {k: v.to(device, dtype=torch.float16) for k, v in inputs.items()} else: inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): if device.type == "cuda": with torch.amp.autocast('cuda'): outputs = model(**inputs) else: outputs = model(**inputs) # Apply softmax to get probabilities logits = outputs.logits probabilities = torch.nn.functional.softmax(logits, dim=-1) probs_np = probabilities.cpu().numpy()[0] # Ensure we have exactly 2 probabilities (AI, Real) if len(probs_np) == 2: # Standard binary classification: [AI, Real] return probs_np else: # Handle edge cases or single-class outputs ai_prob = float(probs_np[0]) if len(probs_np) > 0 else 0.5 real_prob = 1.0 - ai_prob return np.array([ai_prob, real_prob]) except Exception as e: print(f"Prediction error for model: {str(e)}") return None def ensemble_prediction(image): """Enhanced ensemble prediction with improved accuracy""" individual_predictions = [] model_weights = [] model_names = [] start_time = time.time() # Get predictions from each model for model_id, model in models.items(): if model_id in processors: config = MODEL_CONFIGS[model_id] probs = predict_single_model(image, model, processors[model_id], device) if probs is not None and len(probs) >= 2: # Ensure we have [AI_prob, Real_prob] format ai_prob = float(probs[0]) real_prob = float(probs[1]) individual_predictions.append([ai_prob, real_prob]) model_weights.append(config["weight"]) model_names.append(config["name"]) print(f"{config['name']}: AI={ai_prob:.3f}, Real={real_prob:.3f}") analysis_time = time.time() - start_time if not individual_predictions: print("❌ No valid predictions from any model") return None, analysis_time # Weighted ensemble calculation individual_predictions = np.array(individual_predictions) model_weights = np.array(model_weights) # Normalize weights normalized_weights = model_weights / np.sum(model_weights) # Calculate weighted average ensemble_probs = np.average(individual_predictions, axis=0, weights=normalized_weights) ai_prob = float(ensemble_probs[0]) real_prob = float(ensemble_probs[1]) # Ensure probabilities sum to 1.0 total = ai_prob + real_prob if total > 0: ai_prob /= total real_prob /= total print(f"🎯 Ensemble Result: AI={ai_prob:.3f}, Real={real_prob:.3f}") ensemble_result = { "ai_prob": ai_prob, "real_prob": real_prob, "analysis_time": analysis_time, "device": device.type.upper(), "models_used": len(individual_predictions), "confidence": max(ai_prob, real_prob), # Confidence score "prediction": "AI Generated" if ai_prob > real_prob else "Real Image" } return ensemble_result, analysis_time @app.on_event("startup") async def startup_event(): """Load models when the server starts with error handling""" print("🚀 Starting AI vs Real Image Detection API...") try: success = load_models() if success: print("✅ Server startup completed successfully!") else: print("⚠️ Server started but no models loaded!") except Exception as e: print(f"❌ Error during startup: {e}") @app.get("/") async def root(): return { "message": "AI vs Real Image Detection API", "status": "running", "models_loaded": len(models), "available_models": list(MODEL_CONFIGS.keys()) if models else [], "device": str(device) if device else "unknown" } @app.get("/health") async def health_check(): health_status = "healthy" if len(models) > 0 else "degraded" return { "status": health_status, "models_loaded": len(models), "total_models": len(MODEL_CONFIGS), "device": str(device) if device else "unknown", "memory_info": "CPU deployment" } @app.post("/analyze") async def analyze_image(image: UploadFile = File(...)): """Enhanced image analysis with improved accuracy""" # Check if models are loaded if len(models) == 0: raise HTTPException( status_code=503, detail="No models available. Service is starting up or encountered errors." ) # Validate file type if not image.content_type.startswith('image/'): raise HTTPException(status_code=400, detail="File must be an image") try: # Read and process image image_data = await image.read() pil_image = Image.open(io.BytesIO(image_data)) print(f"🔍 Processing image: {image.filename}, size: {len(image_data)} bytes, dimensions: {pil_image.size}") # Analyze image with enhanced ensemble result, analysis_time = ensemble_prediction(pil_image) if result is None: raise HTTPException(status_code=500, detail="Analysis failed - no valid predictions") # Log detailed results confidence_level = "High" if result['confidence'] > 0.8 else "Medium" if result['confidence'] > 0.6 else "Low" print(f"✅ Analysis complete: {result['prediction']} (Confidence: {confidence_level} - {result['confidence']:.3f})") # Return enhanced response enhanced_result = { "prediction": result['prediction'], "ai_prob": round(result['ai_prob'], 4), "real_prob": round(result['real_prob'], 4), "confidence": round(result['confidence'], 4), "confidence_level": confidence_level, "analysis_time": round(result['analysis_time'], 3), "models_used": result['models_used'], "device": result['device'], "image_info": { "filename": image.filename, "size_bytes": len(image_data), "dimensions": pil_image.size, "format": pil_image.format } } return JSONResponse(content=enhanced_result) except Exception as e: print(f"❌ Error analyzing image: {str(e)}") raise HTTPException(status_code=500, detail=f"Analysis error: {str(e)}") if __name__ == "__main__": import os print("🚀 Starting AI vs Real Image Detection API Server...") # Default to port 7860 for Hugging Face Spaces, or use PORT env var port = int(os.getenv("PORT", "7860")) uvicorn.run(app, host="0.0.0.0", port=port)