Update TextGen/router.py

TextGen/router.py

@@ -1,35 +1,74 @@
 from pydantic import BaseModel
-
-from .ConfigEnv import config
+from fastapi import File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-
-from langchain_community.llms import Clarifai
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
+import numpy as np
+from tensorflow.keras.models import load_model
+from PIL import Image
+import io
 from TextGen import app
 
-class Generate(BaseModel):
-    text: str
-
-def generate_text(prompt: str):
-    if prompt == "":
-        return {"detail": "Please provide a prompt."}
-    else:
-        prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
-        llm = Clarifai(
-            pat=config.CLARIFAI_PAT,
-            user_id=config.USER_ID,
-            app_id=config.APP_ID, 
-            model_id=config.MODEL_ID,
-            model_version_id=config.MODEL_VERSION_ID,
-        )
-        llmchain = LLMChain(
-            prompt=prompt,
-            llm=llm
-        )
-        llm_response = llmchain.run({"Prompt": prompt})
-        return Generate(text=llm_response)
+# Response models
+class ImageDetectionResponse(BaseModel):
+    is_ai_generated: bool
+    confidence_score: float
+    confidence_percentage: float
+    prediction_score: float
+    message: str
+
+class ErrorResponse(BaseModel):
+    error: str
+    detail: str
+
+# Global model variable
+model = None
+
+def load_ai_detection_model():
+    """Load the AI detection model"""
+    global model
+    if model is None:
+        try:
+            model = load_model('src/best_model.keras')
+            print("✅ Model loaded successfully")
+        except Exception as e:
+            print(f"❌ Error loading model: {str(e)}")
+            model = None
+    return model
 
+def preprocess_image(image_file):
+    """Preprocess the uploaded image for model prediction"""
+    try:
+        # Read file bytes
+        file_bytes = image_file.read()
+        
+        # Open image using PIL from bytes
+        img = Image.open(io.BytesIO(file_bytes))
+        
+        # Convert to RGB if necessary
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+        
+        # Resize to model's expected input size (300x300)
+        img = img.resize((300, 300), Image.Resampling.LANCZOS)
+        
+        # Convert to array and normalize
+        img_array = np.array(img, dtype=np.float32) / 255.0
+        
+        # Add batch dimension
+        img_array = np.expand_dims(img_array, axis=0)
+        
+        return img_array
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error preprocessing image: {str(e)}")
+
+def predict_image(model, img_array):
+    """Make prediction on the preprocessed image"""
+    try:
+        prediction = model.predict(img_array, verbose=0)
+        return prediction
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error making prediction: {str(e)}")
+
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -40,8 +79,105 @@ app.add_middleware(
 
 @app.get("/", tags=["Home"])
 def api_home():
-    return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
+    return {
+        'message': 'AI Image Detection API',
+        'description': 'Upload an image to detect if it is AI-generated or real',
+        'endpoints': {
+            'POST /detect': 'Upload image for AI detection',
+            'GET /health': 'Check API health status'
+        },
+        'usage': 'Send POST request to /detect with image file',
+        'supported_formats': ['JPG', 'JPEG', 'PNG', 'BMP', 'TIFF']
+    }
 
-@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
-def inference(input_prompt: str):
-    return generate_text(prompt=input_prompt)
+@app.get("/health", tags=["Health"])
+def health_check():
+    """Health check endpoint"""
+    model_status = "loaded" if load_ai_detection_model() is not None else "not_loaded"
+    return {
+        'status': 'healthy',
+        'model_status': model_status,
+        'message': 'AI Image Detection API is running'
+    }
+
+@app.post("/detect", 
+         summary="Detect if image is AI-generated", 
+         tags=["Detection"], 
+         response_model=ImageDetectionResponse,
+         responses={
+             400: {"model": ErrorResponse, "description": "Bad Request"},
+             500: {"model": ErrorResponse, "description": "Internal Server Error"}
+         })
+async def detect_ai_image(file: UploadFile = File(...)):
+    """
+    Upload an image to detect if it's AI-generated or real.
+    
+    - **file**: Image file (JPG, JPEG, PNG, BMP, TIFF)
+    - Returns: Detection result with confidence score
+    """
+    
+    # Validate file type
+    if not file.content_type or not file.content_type.startswith('image/'):
+        raise HTTPException(
+            status_code=400, 
+            detail="Invalid file type. Please upload an image file."
+        )
+    
+    # Check file size (5MB limit)
+    file_size = 0
+    content = await file.read()
+    file_size = len(content)
+    max_size = 5 * 1024 * 1024  # 5MB
+    
+    if file_size > max_size:
+        raise HTTPException(
+            status_code=400,
+            detail=f"File size ({file_size/1024/1024:.2f}MB) exceeds 5MB limit"
+        )
+    
+    # Reset file pointer
+    await file.seek(0)
+    
+    # Load model
+    detection_model = load_ai_detection_model()
+    if detection_model is None:
+        raise HTTPException(
+            status_code=500,
+            detail="AI detection model not available. Please try again later."
+        )
+    
+    try:
+        # Preprocess image
+        img_array = preprocess_image(file.file)
+        
+        # Make prediction
+        prediction = predict_image(detection_model, img_array)
+        
+        # Process results
+        confidence_score = float(prediction[0][0])
+        threshold = 0.5
+        
+        is_ai_generated = confidence_score > threshold
+        
+        if is_ai_generated:
+            confidence_percentage = confidence_score * 100
+            message = "This image appears to be AI-generated"
+        else:
+            confidence_percentage = (1 - confidence_score) * 100
+            message = "This image appears to be real/human-made"
+        
+        return ImageDetectionResponse(
+            is_ai_generated=is_ai_generated,
+            confidence_score=confidence_score,
+            confidence_percentage=round(confidence_percentage, 2),
+            prediction_score=round(confidence_score, 6),
+            message=message
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Unexpected error during image processing: {str(e)}"
+        )