| """ |
| Image classifier for Electrical Outlets. EfficientNet-B0 backbone + MLP head. |
| FINAL v5: 5 classes (no GFCI). |
| """ |
| from pathlib import Path |
| from typing import Dict, Any, Optional |
| import json |
| import torch |
| import torch.nn as nn |
| from torchvision import models |
|
|
|
|
| class ElectricalOutletsImageModel(nn.Module): |
|
|
| def __init__( |
| self, |
| num_classes: int = 5, |
| label_mapping_path: Optional[Path] = None, |
| pretrained: bool = True, |
| head_hidden: int = 256, |
| head_dropout: float = 0.4, |
| ): |
| super().__init__() |
| self.num_classes = num_classes |
| self.backbone = models.efficientnet_b0( |
| weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1 if pretrained else None |
| ) |
| in_features = self.backbone.classifier[1].in_features |
| self.backbone.classifier = nn.Identity() |
|
|
| self.head = nn.Sequential( |
| nn.Dropout(head_dropout), |
| nn.Linear(in_features, head_hidden), |
| nn.ReLU(), |
| nn.Dropout(head_dropout * 0.5), |
| nn.Linear(head_hidden, num_classes), |
| ) |
|
|
| self.idx_to_issue_type = None |
| self.idx_to_severity = None |
| if label_mapping_path and Path(label_mapping_path).exists(): |
| with open(label_mapping_path) as f: |
| lm = json.load(f) |
| self.idx_to_issue_type = lm["image"]["idx_to_issue_type"] |
| self.idx_to_severity = lm["image"]["idx_to_severity"] |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| features = self.backbone(x) |
| return self.head(features) |
|
|
| def predict_to_schema(self, logits: torch.Tensor) -> Dict[str, Any]: |
| probs = torch.softmax(logits, dim=-1) |
| if logits.dim() == 1: |
| probs = probs.unsqueeze(0) |
| conf, pred = probs.max(dim=-1) |
| pred = pred.item() if pred.numel() == 1 else pred |
| conf = conf.item() if conf.numel() == 1 else conf |
| issue_type = (self.idx_to_issue_type or ["unknown"] * self.num_classes)[pred] |
| severity = (self.idx_to_severity or ["medium"] * self.num_classes)[pred] |
| result = "normal" if issue_type == "normal" else "issue_detected" |
| return { |
| "result": result, |
| "issue_type": issue_type, |
| "severity": severity, |
| "confidence": float(conf), |
| "class_idx": int(pred), |
| } |
|
|
