features/real_forged_classifier/controller.py

from typing import IO
import io
import numpy as np
from PIL import Image
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
import torch
from torchvision import transforms

from .preprocessor import preprocessor
from .inferencer import interferencer
from .model_loader import models
from config import Config

security = HTTPBearer()


async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    token = credentials.credentials
    expected_token = Config.SECRET_TOKEN
    if token != expected_token:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid or expired token",
        )
    return token


class ClassificationController:
    """
    Controller to handle the image classification logic.
    """
    def classify_image(self, image_file: IO) -> dict:
        """
        Orchestrates the classification of a single image file.

        Args:
            image_file (IO): The image file to classify.

        Returns:
            dict: The classification result.
        """
        try:
            # Step 1: Preprocess the image
            image_tensor = preprocessor.process(image_file)

            # Step 2: Perform inference
            result = interferencer.predict(image_tensor)

            return result
        except ValueError as e:
            # Handle specific errors like invalid images
            return {"error": str(e)}
        except Exception as e:
            # Handle unexpected errors
            print(f"An unexpected error occurred: {e}")
            return {"error": "An internal error occurred during classification."}

# Create a single instance of the controller
controller = ClassificationController()

class documentForger:
    """
    Document forgery detector that uses the ELA-trained EfficientNet model
    when available (models.doc_model). Returns a dict with verdict and confidence.
    """
    def is_forged(self, document_file: IO) -> dict:
        # Ensure a document model is loaded
        if not hasattr(models, 'doc_model') or models.doc_model is None:
            _downloadmodel = Config.DOCUMENT_FORGERY_MODEL_PATH
            return {"detail": "Document forgery model not available."}

        # Read file bytes
        try:
            data = document_file.read()
            img = Image.open(io.BytesIO(data)).convert('RGB')
        except Exception as e:
            return {"detail": f"Could not open document image: {e}"}

        # Compute ELA map (same approach as the notebook)
        try:
            buf = io.BytesIO()
            img.save(buf, format='JPEG', quality=90)
            buf.seek(0)
            recompressed = Image.open(buf).convert('RGB')

            ela_arr = np.abs(np.array(img, dtype=np.float32) - np.array(recompressed, dtype=np.float32))
            p99 = np.percentile(ela_arr, 99)
            if p99 > 0:
                ela_arr = np.clip(ela_arr * (255.0 / p99), 0, 255).astype(np.uint8)
            else:
                ela_arr = ela_arr.astype(np.uint8)

            ela_pil = Image.fromarray(ela_arr, mode='RGB')
        except Exception as e:
            return {"detail": f"Failed to compute ELA: {e}"}

        # Transform and run through model
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])

        tensor = transform(ela_pil).unsqueeze(0).to(models.device)

        with torch.no_grad():
            logits = models.doc_model(tensor)
            probs = torch.softmax(logits, dim=1)[0, 1].item()

        # Interpret confidence using configurable thresholds (values in 0..1)
        low = getattr(Config, 'DOCUMENT_FORGERY_POSSIBLE_LOW', 0.40)
        high = getattr(Config, 'DOCUMENT_FORGERY_FORGED_LOW', 0.55)

        if probs < low:
            verdict = 'LIKELY AUTHENTIC'
        elif probs < high:
            verdict = 'POSSIBLY FORGED'
        else:
            verdict = 'LIKELY FORGED'

        return {
            "verdict": verdict,
            "confidence": float(probs),
            "confidence_pct": round(float(probs) * 100, 2),
        }

# Create a single instance of the document forger
document_forger = documentForger()