Update main.py

main.py

@@ -1,116 +1,220 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-import tensorflow as tf
-from tensorflow.keras.models import load_model
-from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
-from PIL import Image
-import numpy as np
-import io
-import os
-
-# 1. INITIALIZE FASTAPI
-app = FastAPI(title="Veritas AI Detector API")
-
-# 2. CORS (CRITICAL: Allows your frontend to talk to this backend)
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allows all origins. For production, change to your frontend URL.
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# 3. GLOBAL MODEL LOADER
-# We load the model once at startup so we don't reload it for every request.
-MODEL_PATH = "with_flux_model.keras"
-model = None
-
-@app.on_event("startup")
-async def load_ai_model():
-    global model
-    if os.path.exists(MODEL_PATH):
-        print(f"🔄 Loading Model from {MODEL_PATH}...")
-        # compile=False makes it load faster/safer for inference only
-        model = load_model(MODEL_PATH, compile=False)
-        print(f"✅ Model Loaded Successfully!")
-    else:
-        print(f"❌ CRITICAL ERROR: Model file '{MODEL_PATH}' not found.")
-
-# 4. PREPROCESSING FUNCTION (The "Perfect" Match)
-def preprocess_image(image_bytes):
-    """
-    Exact replica of the training preprocessing.
-    1. Open Image
-    2. Convert to RGB (Fixes PNG/Transparency issues)
-    3. Resize to 224x224 using LANCZOS (High quality)
-    4. Apply EfficientNetV2 specific preprocessing
-    """
-    try:
-        # Open image from bytes
-        img = Image.open(io.BytesIO(image_bytes))
-        
-        # Force RGB
-        img = img.convert('RGB')
-        
-        # Resize (LANCZOS is what we used in training)
-        img = img.resize((224, 224), Image.Resampling.LANCZOS)
-        
-        # Convert to Array
-        img_array = np.array(img)
-        
-        # Expand dims (1, 224, 224, 3)
-        img_array = np.expand_dims(img_array, axis=0)
-        
-        # TF EfficientNet Preprocessing
-        return preprocess_input(img_array)
-        
-    except Exception as e:
-        raise ValueError(f"Image processing failed: {str(e)}")
-
-# 5. THE PREDICTION ENDPOINT
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    global model
-    
-    # Validation
-    if not model:
-        raise HTTPException(status_code=500, detail="Model not loaded")
-    
-    if file.content_type not in ["image/jpeg", "image/png", "image/jpg", "image/webp"]:
-        raise HTTPException(status_code=400, detail="Invalid file type. Please upload JPG or PNG.")
-
-    try:
-        # Read file
-        contents = await file.read()
-        
-        # Preprocess
-        processed_image = preprocess_image(contents)
-        
-        # Predict
-        prediction = model.predict(processed_image)
-        
-        # Classes (Alphabetical Order used by Keras: 0=AI, 1=Real)
-        ai_score = float(prediction[0][0])
-        real_score = float(prediction[0][1])
-        
-        # Logic
-        result_class = "AI" if ai_score > real_score else "Real"
-        confidence = max(ai_score, real_score) * 100
-        
-        return {
-            "prediction": result_class,
-            "confidence_percentage": round(confidence, 2),
-            "probabilities": {
-                "ai": round(ai_score, 4),
-                "real": round(real_score, 4)
-            },
-            "status": "success"
-        }
-
-    except Exception as e:
-        return {"status": "error", "detail": str(e)}
-
-# Health Check
-@app.get("/")
-def home():
-    return {"status": "online", "message": "Veritas AI Detector is running."}
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
+from PIL import Image
+import numpy as np
+import cv2
+import io
+import os
+import base64
+import tempfile
+
+# ==========================================
+# 1. INITIAL SETUP
+# ==========================================
+app = FastAPI(title="Veritas Forensic Engine")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+MODEL_PATH = "with_flux_model.keras"
+model = None
+
+# We need to find the last convolutional layer for GradCAM
+# For EfficientNetV2B0, it is usually 'top_activation'
+LAST_CONV_LAYER_NAME = "top_activation"
+
+@app.on_event("startup")
+async def load_ai_model():
+    global model
+    if os.path.exists(MODEL_PATH):
+        print(f"🔄 Loading Forensic Model...")
+        model = load_model(MODEL_PATH, compile=False)
+        print(f"✅ Model Loaded. Ready for Image & Video Analysis.")
+    else:
+        print(f"❌ CRITICAL: {MODEL_PATH} not found.")
+
+# ==========================================
+# 2. CORE PROCESSING (Sanitization)
+# ==========================================
+def preprocess_image_array(img_array):
+    """
+    Standardizes input for the model.
+    Expects numpy array (RGB). Returns batch tensor.
+    """
+    # Resize using TensorFlow to match training exactly
+    img_resized = tf.image.resize(img_array, (224, 224), method='lanczos3')
+    img_expanded = tf.expand_dims(img_resized, axis=0)
+    return preprocess_input(img_expanded)
+
+def make_gradcam_heatmap(img_array, pred_index=None):
+    """
+    Generates the 'X-Ray' Heatmap using gradients.
+    """
+    # Create a model that maps the input image to the activations of the last conv layer
+    # as well as the output predictions
+    grad_model = tf.keras.models.Model(
+        [model.inputs], 
+        [model.get_layer(LAST_CONV_LAYER_NAME).output, model.output]
+    )
+
+    with tf.GradientTape() as tape:
+        last_conv_layer_output, preds = grad_model(img_array)
+        if pred_index is None:
+            pred_index = tf.argmax(preds[0])
+        class_channel = preds[:, pred_index]
+
+    # Compute the gradient of the output neuron (class) with regard to the output feature map
+    grads = tape.gradient(class_channel, last_conv_layer_output)
+
+    # Average gradients spatially
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+
+    # Weight the output feature map by the importance of the class
+    last_conv_layer_output = last_conv_layer_output[0]
+    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
+    heatmap = tf.squeeze(heatmap)
+
+    # Normalize the heatmap between 0 and 1
+    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
+    return heatmap.numpy()
+
+def overlay_heatmap(original_img_pil, heatmap):
+    """
+    Merges the Heatmap with the original image for display.
+    """
+    # Convert PIL to Array
+    img = np.array(original_img_pil)
+    
+    # Rescale heatmap to a range 0-255
+    heatmap = np.uint8(255 * heatmap)
+
+    # Use jet colormap to colorize heatmap
+    jet = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+
+    # Resize heatmap to match original image size
+    jet = cv2.resize(jet, (img.shape[1], img.shape[0]))
+
+    # Superimpose the heatmap on original image
+    # 0.4 = Heatmap intensity, 0.6 = Original Image intensity
+    superimposed_img = jet * 0.4 + img * 0.6
+    superimposed_img = np.clip(superimposed_img, 0, 255).astype("uint8")
+
+    # Encode back to Base64 to send to Frontend
+    is_success, buffer = cv2.imencode(".jpg", superimposed_img)
+    return base64.b64encode(buffer).decode("utf-8")
+
+# ==========================================
+# 3. ENDPOINT: IMAGE FORENSICS
+# ==========================================
+@app.post("/api/analyze-image")
+async def analyze_image(file: UploadFile = File(...)):
+    if not model: raise HTTPException(500, "Model loading...")
+    
+    try:
+        # 1. Read & Sanitize
+        contents = await file.read()
+        img = Image.open(io.BytesIO(contents)).convert('RGB')
+        
+        # 2. Prepare for Model
+        processed_tensor = preprocess_image_array(np.array(img))
+        
+        # 3. Predict
+        preds = model.predict(processed_tensor)
+        ai_score = float(preds[0][0])
+        real_score = float(preds[0][1])
+        
+        confidence = max(ai_score, real_score) * 100
+        label = "AI" if ai_score > real_score else "Real"
+
+        # 4. Generate Heatmap (Only if AI is suspected, or always if you prefer)
+        # We generate it always so the user can toggle it
+        heatmap = make_gradcam_heatmap(processed_tensor)
+        heatmap_b64 = overlay_heatmap(img, heatmap)
+
+        return {
+            "type": "image",
+            "prediction": label,
+            "confidence": round(confidence, 2),
+            "heatmap_base64": heatmap_b64, # The visual proof
+            "probabilities": {"ai": ai_score, "real": real_score}
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
+
+# ==========================================
+# 4. ENDPOINT: VIDEO SENTINEL
+# ==========================================
+@app.post("/api/analyze-video")
+async def analyze_video(file: UploadFile = File(...)):
+    if not model: raise HTTPException(500, "Model loading...")
+    
+    try:
+        # 1. Save Video Temporarily (OpenCV needs a file path)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_vid:
+            temp_vid.write(await file.read())
+            temp_path = temp_vid.name
+
+        # 2. Process Video
+        cap = cv2.VideoCapture(temp_path)
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        duration = frame_count / fps if fps > 0 else 0
+        
+        # We will extract 10 Keyframes evenly spaced
+        frames_to_analyze = 10
+        step = max(1, frame_count // frames_to_analyze)
+        
+        timeline_results = []
+        fake_frame_count = 0
+        
+        for i in range(0, frame_count, step):
+            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+            ret, frame = cap.read()
+            if not ret: break
+            
+            # Convert BGR (OpenCV) to RGB (Model)
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            
+            # Preprocess & Predict
+            processed_tensor = preprocess_image_array(frame_rgb)
+            preds = model.predict(processed_tensor)
+            ai_score = float(preds[0][0])
+            
+            # Timestamp
+            timestamp = round(i / fps, 2) if fps > 0 else 0
+            
+            timeline_results.append({
+                "timestamp": timestamp,
+                "ai_score": ai_score,
+                "status": "FAKE" if ai_score > 0.5 else "REAL"
+            })
+            
+            if ai_score > 0.5: fake_frame_count += 1
+            
+            if len(timeline_results) >= frames_to_analyze: break
+            
+        cap.release()
+        os.unlink(temp_path) # Delete temp file
+
+        # 3. Final Video Verdict
+        overall_fake_percent = (fake_frame_count / len(timeline_results)) * 100
+        final_verdict = "DEEPFAKE DETECTED" if overall_fake_percent > 40 else "AUTHENTIC VIDEO"
+
+        return {
+            "type": "video",
+            "prediction": final_verdict,
+            "fake_percentage": round(overall_fake_percent, 2),
+            "timeline": timeline_results # List of frame analysis
+        }
+
+    except Exception as e:
+        return {"error": str(e)}