Spaces:
Running
Running
| """DHVANI proprietary custom model runner.""" | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from core.config import SAMPLE_RATE, custom_model_path | |
| from core.dhvani_head import BACKBONE_ID, DhvaniCustomModel, MAX_SAMPLES | |
| logger = logging.getLogger("dhvani.custom") | |
| class CustomDhvaniDetector: | |
| """Loads models/dhvani-custom.pt trained via training/train_dhvani_head.py.""" | |
| def __init__(self, checkpoint_path: str | None = None, device: str | None = None) -> None: | |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") | |
| self.checkpoint_path = checkpoint_path or custom_model_path() | |
| self._model: DhvaniCustomModel | None = None | |
| self._backbone_id = BACKBONE_ID | |
| self._fake_label_index = 1 | |
| self._metadata: dict = {} | |
| def enabled(self) -> bool: | |
| return bool(self.checkpoint_path) and Path(self.checkpoint_path).exists() | |
| def load(self) -> None: | |
| if not self.enabled: | |
| return | |
| if self._model is not None: | |
| return | |
| logger.info("Loading DHVANI custom model from %s", self.checkpoint_path) | |
| payload = torch.load(self.checkpoint_path, map_location="cpu", weights_only=False) | |
| self._backbone_id = payload.get("backbone_id", self._backbone_id) | |
| self._fake_label_index = int(payload.get("fake_label_index", 1)) | |
| self._metadata = {k: v for k, v in payload.items() if k != "head_state_dict"} | |
| self._model = DhvaniCustomModel(backbone_id=self._backbone_id) | |
| self._model.head.load_state_dict(payload["head_state_dict"]) | |
| self._model.to(self.device) | |
| self._model.eval() | |
| def loaded(self) -> bool: | |
| return self._model is not None | |
| def model_id(self) -> str: | |
| version = self._metadata.get("version", "") | |
| if version: | |
| return str(version) | |
| return "dhvani-custom-head" | |
| def version(self) -> str: | |
| return str(self._metadata.get("version", "unknown")) | |
| def metadata(self) -> dict: | |
| return dict(self._metadata) | |
| def fake_probability(self, waveform: np.ndarray) -> float: | |
| if not self.enabled: | |
| return 0.0 | |
| self.load() | |
| assert self._model is not None | |
| if waveform.size > MAX_SAMPLES: | |
| waveform = waveform[:MAX_SAMPLES] | |
| if waveform.size < SAMPLE_RATE: | |
| waveform = np.pad(waveform, (0, SAMPLE_RATE - waveform.size)) | |
| tensor = torch.from_numpy(waveform.astype(np.float32)).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| logits = self._model(tensor) | |
| probs = torch.softmax(logits, dim=-1).cpu().numpy()[0] | |
| return float(probs[self._fake_label_index]) |