features/real_forged_classifier/controller.py

from typing import IO
import io
import numpy as np
from PIL import Image
import torch
from torchvision import transforms

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


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:
            return {"error": "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 {"error": 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 {"error": 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()