File size: 2,773 Bytes
998740d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """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) # (B, P, d)
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"]
# Normalisation statistics from training (uint8/255 scale)
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
|