from __future__ import annotations import logging import time from pathlib import Path from typing import Any import torch from PIL import Image from torchvision import models, transforms from app.config import settings logger = logging.getLogger(__name__) class ClassifierService: def __init__(self) -> None: self.model: torch.nn.Module | None = None self.model_name = "bitcheck_efficientnet_b0" self.model_status = "not_loaded" self.warning: str | None = None self.device = torch.device("cpu") self.threshold = settings.model_threshold self.image_size = settings.model_input_size self.normalization_mean = [0.485, 0.456, 0.406] self.normalization_std = [0.229, 0.224, 0.225] self.id2label = {0: "real", 1: "ai_generated"} self.transform = self._build_transform() self.model_path: Path | None = None @property def is_loaded(self) -> bool: return self.model is not None and self.model_status == "loaded" def load(self) -> None: if self.model is not None: return path = self._resolve_model_path() if path is None: self.model_status = "missing" self.warning = f"No model file found at {settings.model_path}." return try: checkpoint = torch.load(path, map_location=self.device) state_dict = checkpoint.get("model_state_dict", checkpoint.get("state_dict", checkpoint)) if isinstance(checkpoint, dict) else checkpoint self._apply_checkpoint_metadata(checkpoint if isinstance(checkpoint, dict) else {}) self.model = self._build_model(state_dict) self.model.load_state_dict(state_dict, strict=False) self.model.to(self.device) self.model.eval() self.model_path = path self.model_status = "loaded" self.warning = None if path == settings.model_path else f"Loaded fallback model file: {path.name}." except Exception as exc: logger.exception("PyTorch classifier failed to load") self.model = None self.model_status = "error" self.warning = str(exc) def predict_path(self, image_path: Path, threshold: float | None = None) -> dict[str, Any]: with Image.open(image_path) as img: return self.predict(img.convert("RGB"), threshold=threshold) def predict(self, image: Image.Image, threshold: float | None = None) -> dict[str, Any]: start = time.perf_counter() if self.model is None: self.load() if self.model is None: return { "checked": False, "model_status": self.model_status, "model_name": self.model_name, "predicted_label": "unknown", "real_probability": None, "ai_generated_probability": None, "threshold": threshold if threshold is not None else self.threshold, "risk_score": None, "inference_time_ms": round((time.perf_counter() - start) * 1000, 2), "warning": self.warning or "Classifier model is not available.", } try: active_threshold = float(threshold if threshold is not None else self.threshold) tensor = self.transform(image.convert("RGB")).unsqueeze(0).to(self.device) with torch.no_grad(): logits = self.model(tensor) probs = torch.softmax(logits, dim=1)[0].detach().cpu().tolist() real_prob = float(probs[0]) if len(probs) > 0 else 0.0 ai_prob = float(probs[1]) if len(probs) > 1 else 1.0 - real_prob predicted_index = 1 if ai_prob >= active_threshold else 0 predicted_label = "likely_ai_generated" if predicted_index == 1 else "likely_authentic" return { "checked": True, "model_status": "loaded", "model_name": self.model_name, "predicted_label": predicted_label, "real_probability": round(real_prob, 4), "ai_generated_probability": round(ai_prob, 4), "threshold": active_threshold, "risk_score": round(ai_prob, 4), "inference_time_ms": round((time.perf_counter() - start) * 1000, 2), "warning": "Classifier output is probabilistic and may not generalize to unseen generators.", } except Exception as exc: return { "checked": False, "model_status": "error", "model_name": self.model_name, "predicted_label": "unknown", "real_probability": None, "ai_generated_probability": None, "threshold": threshold if threshold is not None else self.threshold, "risk_score": None, "inference_time_ms": round((time.perf_counter() - start) * 1000, 2), "warning": "Classifier inference failed.", "error": str(exc), } def _resolve_model_path(self) -> Path | None: if settings.model_path.exists(): return settings.model_path for path in settings.model_fallback_paths: if path.exists(): return path return None def _apply_checkpoint_metadata(self, checkpoint: dict[str, Any]) -> None: self.model_name = str(checkpoint.get("model_name") or checkpoint.get("architecture") or self.model_name) self.image_size = int(checkpoint.get("image_size") or checkpoint.get("img_size") or self.image_size) self.threshold = float(checkpoint.get("threshold") or self.threshold) self.normalization_mean = list(checkpoint.get("normalization_mean") or self.normalization_mean) self.normalization_std = list(checkpoint.get("normalization_std") or self.normalization_std) class_names = checkpoint.get("class_names") if isinstance(class_names, list) and len(class_names) >= 2: self.id2label = {idx: str(label) for idx, label in enumerate(class_names)} if isinstance(checkpoint.get("id2label"), dict): self.id2label = {int(k): str(v) for k, v in checkpoint["id2label"].items()} self.transform = self._build_transform() def _build_transform(self) -> transforms.Compose: return transforms.Compose( [ transforms.Resize((self.image_size, self.image_size)), transforms.ToTensor(), transforms.Normalize(mean=self.normalization_mean, std=self.normalization_std), ] ) def _build_model(self, state_dict: dict[str, torch.Tensor]) -> torch.nn.Module: out_features = 2 weight = state_dict.get("classifier.1.weight") if isinstance(weight, torch.Tensor): out_features = int(weight.shape[0]) model = models.efficientnet_b0(weights=None) in_features = model.classifier[1].in_features model.classifier[1] = torch.nn.Linear(in_features, out_features) return model def target_layer(self) -> torch.nn.Module | None: if self.model is None: return None last_conv = None for module in self.model.modules(): if isinstance(module, torch.nn.Conv2d): last_conv = module return last_conv classifier_service = ClassifierService()