Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| from PIL import Image | |
| from core.ml_utils import ( | |
| batched, | |
| heatmap_from_occlusion_scores, | |
| labels_from_config, | |
| model_heatmap_batch_size, | |
| model_heatmap_grid_size, | |
| model_heatmaps_enabled, | |
| occluded_images, | |
| prepare_rgb, | |
| preferred_device, | |
| probability_from_labels, | |
| softmax, | |
| ) | |
| from detectors.base import BaseDetector, DetectorResult | |
| class EfficientNetDetector(BaseDetector): | |
| name = "EfficientNet Detector" | |
| category = "ml_model" | |
| weight = 1.25 | |
| default_model = "KakoSan/siftai-efficientnet-b4-v2" | |
| _model = None | |
| _processor = None | |
| _device = None | |
| _loaded_model_id = None | |
| def _load(self): | |
| model_id = os.getenv("EFFICIENTNET_FORENSICS_MODEL", self.default_model) | |
| if self._model is not None and self._loaded_model_id == model_id: | |
| return model_id, self._processor, self._model, self._device | |
| try: | |
| import torch | |
| from transformers import AutoImageProcessor, AutoModelForImageClassification | |
| except Exception as exc: | |
| raise RuntimeError(f"torch/transformers are not installed: {exc}") from exc | |
| processor = AutoImageProcessor.from_pretrained(model_id) | |
| model = AutoModelForImageClassification.from_pretrained(model_id) | |
| device = preferred_device() | |
| model.to(device) | |
| model.eval() | |
| self.__class__._processor = processor | |
| self.__class__._model = model | |
| self.__class__._device = device | |
| self.__class__._loaded_model_id = model_id | |
| return model_id, processor, model, device | |
| def analyze(self, image: Image.Image, context: dict) -> DetectorResult: | |
| try: | |
| model_id, processor, model, device = self._load() | |
| import torch | |
| inputs = processor(images=prepare_rgb(image), return_tensors="pt").to(device) | |
| with torch.inference_mode(): | |
| outputs = model(**inputs) | |
| probs = softmax(outputs.logits.detach().float().cpu().numpy()[0]) | |
| labels = labels_from_config(model, len(probs)) | |
| fake_probability, top_score, top_label = probability_from_labels(labels, probs) | |
| artifacts = {"model_id": model_id, "device": device, "labels": labels} | |
| if model_heatmaps_enabled(): | |
| grid_size = model_heatmap_grid_size() | |
| occluded = occluded_images(image, grid_size) | |
| occluded_scores: list[float] = [] | |
| for batch in batched(occluded, model_heatmap_batch_size()): | |
| batch_inputs = processor(images=[prepare_rgb(item) for item in batch], return_tensors="pt").to(device) | |
| with torch.inference_mode(): | |
| batch_outputs = model(**batch_inputs) | |
| batch_probs = batch_outputs.logits.detach().float().cpu().numpy() | |
| for logits in batch_probs: | |
| prob, _, _ = probability_from_labels(labels, softmax(logits)) | |
| occluded_scores.append(prob) | |
| artifacts["efficientnet_heatmap"] = heatmap_from_occlusion_scores( | |
| image, fake_probability, occluded_scores, grid_size | |
| ) | |
| artifacts["heatmap_method"] = f"{grid_size}x{grid_size} occlusion sensitivity" | |
| confidence = max(0.35, min(0.92, top_score)) | |
| explanation = ( | |
| f"Loaded EfficientNet-family image-classification model `{model_id}` on {device}; " | |
| f"top label `{top_label}` with score {top_score:.3f}." | |
| ) | |
| return self.result( | |
| ai_probability=0.38 + fake_probability * 0.45, | |
| manipulation_probability=fake_probability, | |
| confidence=confidence, | |
| explanation=explanation, | |
| warnings=["EfficientNet deepfake models may be strongest on faces or training-domain imagery."], | |
| artifacts=artifacts, | |
| ) | |
| except RuntimeError as exc: | |
| return self.not_loaded_result(str(exc)) | |
| except Exception as exc: | |
| return self.error_result(exc) | |