| """TinyViT for surface-defect image classification (standalone, no project deps).""" |
| import torch |
| import torch.nn as nn |
|
|
|
|
| class TinyViT(nn.Module): |
| """Image-classification ViT: img_size x img_size -> patch x patch patches.""" |
|
|
| def __init__(self, img_size: int = 64, patch: int = 8, d_model: int = 64, |
| n_layers: int = 2, n_heads: int = 4, n_classes: int = 6, |
| in_ch: int = 1, dropout: float = 0.1): |
| super().__init__() |
| assert img_size % patch == 0, "img_size must be divisible by patch" |
| self.img_size = img_size |
| self.patch = patch |
| self.n_patches = (img_size // patch) ** 2 |
| self.patch_embed = nn.Conv2d(in_ch, d_model, |
| kernel_size=patch, stride=patch) |
| self.cls = nn.Parameter(torch.zeros(1, 1, d_model)) |
| self.pos = nn.Parameter( |
| torch.zeros(1, self.n_patches + 1, d_model)) |
| layer = nn.TransformerEncoderLayer( |
| d_model=d_model, nhead=n_heads, |
| dim_feedforward=4 * d_model, dropout=dropout, |
| batch_first=True, activation="gelu") |
| self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers) |
| self.norm = nn.LayerNorm(d_model) |
| self.head = nn.Sequential( |
| nn.Dropout(dropout), |
| nn.Linear(d_model, d_model), |
| nn.GELU(), |
| nn.Linear(d_model, n_classes)) |
|
|
| def forward(self, x): |
| B = x.shape[0] |
| t = self.patch_embed(x).flatten(2).transpose(1, 2) |
| t = torch.cat([self.cls.expand(B, -1, -1), t], dim=1) + self.pos |
| t = self.encoder(t) |
| return self.head(self.norm(t)[:, 0]) |
|
|
|
|
| CLASS_NAMES = ["crazing", "inclusion", "patches", |
| "pitted_surface", "rolled-in_scale", "scratches"] |
|
|
| |
| MEAN = 0.5013 |
| STD = 0.2075 |
|
|
|
|
| def load_model(weights_path: str, device: str = "cpu") -> TinyViT: |
| model = TinyViT() |
| state = torch.load(weights_path, map_location=device) |
| model.load_state_dict(state) |
| model.eval() |
| return model |
|
|
|
|
| def predict_image(model: TinyViT, image) -> tuple[int, "np.ndarray"]: |
| """Classify a PIL image or numpy array (grayscale or RGB). |
| |
| Returns (predicted_class_id, class_probabilities[6]). |
| """ |
| import numpy as np |
| from PIL import Image |
| if isinstance(image, np.ndarray): |
| image = Image.fromarray(image) |
| image = image.convert("L").resize((64, 64), Image.BILINEAR) |
| x = np.asarray(image, dtype=np.float32) / 255.0 |
| x = (x - MEAN) / STD |
| xt = torch.from_numpy(x)[None, None].to(next(model.parameters()).device) |
| with torch.no_grad(): |
| probs = torch.softmax(model(xt), dim=1).cpu().numpy()[0] |
| return int(probs.argmax()), probs |
|
|
