File size: 3,801 Bytes
6396e6b c8f4e3f 6396e6b c8f4e3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | 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()
|