Spaces:
Sleeping
Sleeping
Hikity
feat: add Pipeline 3 scene detection (YOLO road damage + TrashNet) and DistilBERT priority validator for dual-model consensus routing
39f212a | import os | |
| import re | |
| from pathlib import Path | |
| from functools import partial | |
| import torch | |
| from torch import nn | |
| from torch.nn.modules.dropout import Dropout | |
| from torch.nn.modules.linear import Linear | |
| from torch.nn.modules.pooling import AdaptiveAvgPool2d | |
| try: | |
| from timm.models.efficientnet import tf_efficientnet_b7_ns | |
| except ImportError: | |
| def tf_efficientnet_b7_ns(*args, **kwargs): | |
| return nn.Sequential() | |
| encoder_params = { | |
| "tf_efficientnet_b7_ns": { | |
| "features": 2560, | |
| "init_op": partial(tf_efficientnet_b7_ns, pretrained=False, drop_path_rate=0.2) | |
| } | |
| } | |
| class DeepFakeClassifier(nn.Module): | |
| def __init__(self, encoder="tf_efficientnet_b7_ns", dropout_rate=0.0) -> None: | |
| super().__init__() | |
| self.encoder = encoder_params[encoder]["init_op"]() | |
| self.avg_pool = AdaptiveAvgPool2d((1, 1)) | |
| self.dropout = Dropout(dropout_rate) | |
| self.fc = Linear(encoder_params[encoder]["features"], 1) | |
| def forward(self, x): | |
| x = self.encoder.forward_features(x) | |
| x = self.avg_pool(x).flatten(1) | |
| x = self.dropout(x) | |
| x = self.fc(x) | |
| return x | |
| def load_models(weights_dir: str, device: str) -> list: | |
| """Loads all .pth files from weights/ directory.""" | |
| models = [] | |
| weights_path = Path(weights_dir) | |
| if not weights_path.exists(): | |
| print(f"Weights directory {weights_dir} does not exist.") | |
| return models | |
| # Check for weight files matching 'final_*' | |
| weight_files = sorted(list(weights_path.glob("final_*"))) | |
| if not weight_files: | |
| print(f"No checkpoint files starting with 'final_' found in {weights_dir}") | |
| # Look for general .pth or .pt files if final_ is not found | |
| weight_files = sorted(list(weights_path.glob("*.pth"))) + sorted(list(weights_path.glob("*.pt"))) | |
| for weight_file in weight_files: | |
| print(f"Loading weights from {weight_file}") | |
| model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns").to(device) | |
| try: | |
| try: | |
| checkpoint = torch.load(weight_file, map_location="cpu", weights_only=False) | |
| except TypeError: | |
| checkpoint = torch.load(weight_file, map_location="cpu") | |
| state_dict = checkpoint.get("state_dict", checkpoint) | |
| # Strip "module." prefix if model was trained with DataParallel | |
| clean_state_dict = {re.sub("^module.", "", k): v for k, v in state_dict.items()} | |
| model.load_state_dict(clean_state_dict, strict=True) | |
| model.eval() | |
| if device == "cuda": | |
| models.append(model.half()) # FP16 for speed on GPU | |
| else: | |
| models.append(model.float()) # FP32 for CPU compatibility | |
| print(f"Successfully loaded checkpoint {weight_file.name}") | |
| except Exception as e: | |
| print(f"Error loading checkpoint {weight_file.name}: {e}") | |
| return models | |