Spaces:
Running
Running
| 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 | |
| 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() | |