Views
No views yet
torchvision.models.mobilenet_v3_small (ImageNet-pretrained, IMAGENET1K_V1)nn.Linear in classifier with nn.Linear(in_features, 1)torch.sigmoid for P(corrupted)0.51from PIL import Image
2
3def letterbox(img: Image.Image, size: int = 224) -> Image.Image:
4 w, h = img.size
5 scale = size / max(w, h)
6 new_w, new_h = int(w * scale), int(h * scale)
7 img = img.resize((new_w, new_h), Image.BILINEAR)
8 padded = Image.new("RGB", (size, size), (0, 0, 0))
9 padded.paste(img, ((size - new_w) // 2, (size - new_h) // 2))
10 return padded1import torch
2from torch import nn
3from torchvision import transforms
4from torchvision.models import mobilenet_v3_small
5from huggingface_hub import hf_hub_download
6from safetensors.torch import load_file
7from PIL import Image
8
9weights_path = hf_hub_download(
10 repo_id="callum-sh/codec-corruption-classifier",
11 filename="model.safetensors",
12)
13state = load_file(weights_path)
14
15model = mobilenet_v3_small(weights=None)
16in_features = model.classifier[-1].in_features
17model.classifier[-1] = nn.Linear(in_features, 1)
18model.load_state_dict(state)
19model.eval()
20
21tx = transforms.Compose([
22 transforms.Lambda(lambda im: letterbox(im, 224)),
23 transforms.ToTensor(),
24 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
25])
26
27img = Image.open("frame.jpg").convert("RGB")
28with torch.no_grad():
29 logit = model(tx(img).unsqueeze(0))
30 p_corrupted = torch.sigmoid(logit).item()
31print(f"P(corrupted) = {p_corrupted:.3f}")P(corrupted) > 0.5 should be excluded from the SfM input set.