app.py

import gradio as gr
import os
import cv2
import tempfile
import warnings
from PIL import Image
import numpy as np
import time
import re
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoImageProcessor, AutoModelForImageClassification
from typing import Dict

# --- Custom Tools ---
try:
    from custom_tools.video_deepfake_tool import VideoDeepfakeTool
    from custom_tools.image_deepfake_tool import ImageDeepfakeTool
    from custom_tools.liveness_detection_tool import LivenessDetectionTool
    from custom_tools.image_description_tool import ImageDescriptionTool
    from custom_tools.facesimilarity_tool import FaceSimilarityTool
    CUSTOM_TOOLS_AVAILABLE = True
except ImportError as e:
    print(f"ERROR: Could not import one or more custom tools: {e}")
    print("Please ensure custom_tools are in PYTHONPATH and correctly defined.")
    CUSTOM_TOOLS_AVAILABLE = False
    # Define dummy classes
    class VideoDeepfakeTool: 
        def __init__(self): pass
        def apply(self, p): return {"score": 1.0}
    class ImageDeepfakeTool: 
        def __init__(self): pass
        def apply(self, p): return {"label": "Error", "error_message": "Tool not loaded"}
    class LivenessDetectionTool: 
        def __init__(self): pass
        def apply(self, p): return "Error: Tool not loaded"
    class ImageDescriptionTool: 
        def __init__(self): pass
        def apply(self, p): return "Error: Tool not loaded"
    class FaceSimilarityTool: 
        def __init__(self): pass
        def apply(self, p1, p2): return "Error: Tool not loaded"

# Filter warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor*")
warnings.filterwarnings("ignore", message="The attention mask is not set")
warnings.filterwarnings("ignore", message="Setting `pad_token_id` to `eos_token_id`")
warnings.filterwarnings("ignore", message="Truncation was not explicitly activated")

# Thresholds for KYC checks
LIVENESS_THRESHOLD = 0.70
FACE_MATCH_THRESHOLD = 0.05
DEEPFAKE_THRESHOLD = 0.50
MODEL_VERSION = "1.0.0"

# Initialize tools
if CUSTOM_TOOLS_AVAILABLE:
    try:
        #video_deepfake_tool_instance = VideoDeepfakeTool()
        image_deepfake_tool_instance = ImageDeepfakeTool()
        liveness_detection_tool_instance = LivenessDetectionTool()
        #image_description_tool_instance = ImageDescriptionTool()
        face_similarity_tool_instance = FaceSimilarityTool()
        TOOLS_INITIALIZED_SUCCESSFULLY = True
        print("All custom tools initialized successfully.")
    except Exception as e:
        print(f"FATAL ERROR: Failed to initialize one or more tools: {e}")
        TOOLS_INITIALIZED_SUCCESSFULLY = False
else:
    TOOLS_INITIALIZED_SUCCESSFULLY = False

# FastAPI app initialization
app = FastAPI()

# Add CORS middleware to allow cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# KYC Verification API Endpoint
@app.post("/api/kyc-verification/")
async def kyc_verification(user_image: UploadFile = File(...), id_card_image: UploadFile = File(...)):
    """
    API endpoint for KYC verification.
    Accepts user image and ID card image as input and returns verification results.
    """
    try:
        # Save uploaded files temporarily
        user_image_path = f"temp_{user_image.filename}"
        id_card_image_path = f"temp_{id_card_image.filename}"

        with open(user_image_path, "wb") as f:
            f.write(await user_image.read())
        with open(id_card_image_path, "wb") as f:
            f.write(await id_card_image.read())

        # Load images using PIL
        user_image_pil = Image.open(user_image_path)
        id_card_image_pil = Image.open(id_card_image_path)

        # Call the existing process_kyc_verification function
        result = process_kyc_verification(user_image_pil, id_card_image_pil)

        # Clean up temporary files
        os.remove(user_image_path)
        os.remove(id_card_image_path)

        return JSONResponse(content=result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing KYC: {str(e)}")


# Health Check Endpoint
@app.get("/api/health/")
async def health_check():
    """
    API endpoint to check the health of the application.
    """
    return {"status": "healthy", "tools_initialized": TOOLS_INITIALIZED_SUCCESSFULLY}

def process_kyc_verification(user_image: Image.Image, id_card_image: Image.Image):
    """
    Main KYC verification function that processes both images and returns comprehensive results
    
    Args:
        user_image: PIL Image of the user's face photo
        id_card_image: PIL Image of the ID card/document
    
    Returns:
        dict: Comprehensive verification results with scores and analysis
    """
    start_time = time.time()
    
    try:
        # Validate inputs
        if user_image is None:
            return {
                "error": "User image is required",
                "liveness_score": 0.0,
                "face_similarity": 0.0,
                "deepfake_score": 1.0,
                "overall_decision": "REJECTED",
                "processing_time": 0.0
            }
        
        if id_card_image is None:
            return {
                "error": "ID card image is required", 
                "liveness_score": 0.0,
                "face_similarity": 0.0,
                "deepfake_score": 1.0,
                "overall_decision": "REJECTED",
                "processing_time": 0.0
            }
        
        if not TOOLS_INITIALIZED_SUCCESSFULLY:
            return {
                "error": "KYC tools failed to initialize",
                "liveness_score": 0.0,
                "face_similarity": 0.0,
                "deepfake_score": 1.0,
                "overall_decision": "ERROR",
                "processing_time": 0.0
            }
        
        print("Starting KYC verification process...")
        
        # Initialize scores
        liveness_score = 0.0
        face_similarity = 0.0
        deepfake_score = 1.0
        
        # Save images to temporary files for tool processing
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_user:
            user_image.save(temp_user.name, format='JPEG')
            user_image_path = temp_user.name
        
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_id:
            id_card_image.save(temp_id.name, format='JPEG')
            id_card_path = temp_id.name
        
        temp_files_to_clean = [user_image_path, id_card_path]
        
        try:            
            # 2. Liveness Detection
            print("Performing liveness detection...")
            try:
                liveness_result_str = liveness_detection_tool_instance.apply(user_image_path)
                confidence_match = re.search(r"Confidence:\s*([0-9.]+)", liveness_result_str, re.IGNORECASE)
                if confidence_match:
                    liveness_score = float(confidence_match.group(1))
                else:
                    liveness_score = 0.85 if "real" in liveness_result_str.lower() or "live" in liveness_result_str.lower() else 0.25
                liveness_details = liveness_result_str
            except Exception as e:
                liveness_score = 0.0
                liveness_details = f"Error in liveness detection: {str(e)}"
            
            # 3. Face Comparison
            print("Comparing faces...")
            try:
                face_sim_result_str = face_similarity_tool_instance.apply(user_image_path, id_card_path)
                score_match = re.search(r"similarity score:\s*([0-9.]+)", face_sim_result_str, re.IGNORECASE)
                if score_match:
                    face_similarity = float(score_match.group(1))
                else:
                    try:
                        potential_score_str = face_sim_result_str.split(":")[-1].strip().split(" ")[0].replace(')', '').replace('(', '')
                        face_similarity = float(re.sub(r"[^0-9.]", "", potential_score_str))
                    except:
                        face_similarity = 0.075 if "match" in face_sim_result_str.lower() or "true" in face_sim_result_str.lower() else 0.035
                face_details = face_sim_result_str
            except Exception as e:
                face_similarity = 0.0
                face_details = f"Error in face similarity: {str(e)}"
            
            # 4. Deepfake Detection
            print("Detecting deepfakes...")
            try:
                img_df_result = image_deepfake_tool_instance.apply(user_image_path)
                if img_df_result.get("label", "").lower() == "error":
                    deepfake_score = 1.0
                    deepfake_details = img_df_result.get('error_message', 'Unknown error during deepfake detection.')
                else:
                    label = img_df_result.get("label", "AI").lower()
                    label = "ai"
                    confidence = img_df_result.get("confidence", 0.0)
                    if label == "human":
                        deepfake_score = 1.0 - confidence
                    elif label == "ai":
                        deepfake_score = 1.0 #confidence
                    else:
                        deepfake_score = 1.0
                    deepfake_details = f"Prediction: '{label.capitalize()}' with confidence {confidence:.4f}"
            except Exception as e:
                deepfake_score = 1.0
                deepfake_details = f"Error in deepfake detection: {str(e)}"
            
            # 5. Overall Decision Logic
            processing_time = time.time() - start_time
            
            # Decision criteria
            liveness_pass = liveness_score >= LIVENESS_THRESHOLD
            face_match_pass = face_similarity >= FACE_MATCH_THRESHOLD
            deepfake_pass = deepfake_score <= DEEPFAKE_THRESHOLD
            
            overall_decision = "APPROVED" if (liveness_pass and face_match_pass and deepfake_pass) else "REJECTED"
            confidence_overall = min(liveness_score, face_similarity, 1 - deepfake_score)
            
            # Generate recommendations
            recommendations = []
            if liveness_pass:
                recommendations.append("✅ Liveness check passed - Live person detected")
            else:
                recommendations.append("❌ Liveness check failed - Manual review required")
                
            if face_match_pass:
                recommendations.append("✅ Face similarity acceptable - Identity match confirmed")
            else:
                recommendations.append("❌ Face similarity below threshold - Identity mismatch")
                
            if deepfake_pass:
                recommendations.append("✅ No deepfake detected - Image appears authentic")
            else:
                recommendations.append("❌ Potential manipulation detected - Manual verification needed")
            
            # Compile comprehensive results
            result = {
                # Main scores
                "liveness_score": round(float(liveness_score), 3),
                "deepfake_score": round(float(deepfake_score), 3),
                
                # Overall assessment
                "overall_decision": overall_decision,
                
                # Individual check results
                "liveness_passed": liveness_pass,
                "face_match_passed": face_match_pass,
                "deepfake_detected": "Yes",#deepfake_pass,
                
                # Recommendations and metadata
                "recommendations": recommendations,
                "processing_time": round(processing_time, 2),
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()),
                
            }
            
            print(f"KYC verification completed in {processing_time:.2f}s - Decision: {overall_decision}")
            return result
            
        finally:
            # Clean up temporary files
            for temp_path in temp_files_to_clean:
                try:
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
                except Exception as e_clean:
                    print(f"Error cleaning up temp file {temp_path}: {e_clean}")
        
    except Exception as e:
        error_msg = f"KYC verification failed: {str(e)}"
        print(error_msg)
        
        return {
            "error": error_msg,
            "liveness_score": 0.0,
            "face_similarity": 0.0,
            "deepfake_score": 1.0,
            "overall_decision": "ERROR",
            "confidence_overall": 0.0,
            "analysis_details": {"error": str(e)},
            "recommendations": ["Manual review required due to processing error"],
            "processing_time": time.time() - start_time,
            "model_version": MODEL_VERSION,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())
        }

# Create the Gradio interface
def create_interface():
    """Create and configure the Gradio interface"""
    # Custom CSS for better styling
    css = """
    .gradio-container {
        max-width: 1200px !important;
        margin: auto !important;
    }
    .output-json {
        max-height: 500px;
        overflow-y: auto;
    }
    """
    
    # Interface description
    description = """
    Upload a **user photo** and **ID document** for comprehensive AI-powered identity verification.
    """
    
    # Create the interface
    iface = gr.Interface(
        fn=process_kyc_verification,
        inputs=[
            gr.Image(
                type="pil",
                label="👤 User Photo - Upload a clear photo of the person's face"
            ),
            gr.Image(
                type="pil", 
                label="🆔 ID Card/Document - Upload a photo of the government-issued ID document"
            )
        ],
        outputs=gr.JSON(
            label="📊 Verification Results"
        ),
        title="🔒 ReagVis KYC Identity Verification",
        description=description,
        css=css,
        allow_flagging="never",
    )
    return iface


# Launch the application
if __name__ == "__main__":
    import uvicorn

    # Create example files if they don't exist
    example_files_data = {"example_selfie.jpg": (256, 256, 3), "example_id.jpg": (300, 200, 3)}
    current_dir = os.path.dirname(os.path.abspath(__file__)) if "__file__" in locals() else "."
    for fname, shape_data in example_files_data.items():
        fpath = os.path.join(current_dir, fname)
        if not os.path.exists(fpath) and shape_data:
            try:
                Image.fromarray(np.random.randint(0, 256, shape_data, dtype=np.uint8)).save(fpath)
                print(f"Created dummy example file: {fpath}")
            except Exception as e:
                print(f"Could not create dummy image {fpath}: {e}")

    # Launch Gradio interface
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",  # Allow external access
        server_port=7860,       # Default Gradio port
        share=False,            # Set to True for public sharing during development
        debug=False,            # Set to True for debugging
        show_error=True,        # Show detailed errors
        quiet=False             # Set to True to reduce logging
    )

    # Launch FastAPI app
    uvicorn.run(app, host="0.0.0.0", port=8000)