| """Two-page spread vs. single page classifier for document scans. |
| |
| python inference.py scan1.jpg scan2.jpg ... |
| |
| Prints the logit and the decision for each scan. |
| """ |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from PIL import Image, ImageOps |
| from safetensors.torch import load_file |
| from torchvision.models import efficientnet_b0 |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| class SpreadNet(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.features = efficientnet_b0(weights=None).features |
| self.cls = nn.Sequential(nn.Dropout(0.3), nn.Linear(1280, 1)) |
|
|
| def forward(self, x): |
| return self.cls(self.features(x).mean((2, 3))).squeeze(1) |
|
|
|
|
| def load_model(folder=HERE): |
| cfg = json.loads((Path(folder) / "config.json").read_text()) |
| model = SpreadNet() |
| model.load_state_dict(load_file(str(Path(folder) / "model.safetensors"))) |
| return model.eval(), cfg |
|
|
|
|
| def preprocess(path, cfg): |
| """RGB, longest side = image_size, centred on a gray square, ImageNet normalisation.""" |
| size = cfg["image_size"] |
| im = ImageOps.exif_transpose(Image.open(path)).convert("RGB") |
| k = size / max(im.size) |
| im = im.resize((max(1, round(im.width * k)), max(1, round(im.height * k))), Image.LANCZOS) |
| sq = Image.new("RGB", (size, size), (cfg["pad_value"],) * 3) |
| sq.paste(im, ((size - im.width) // 2, (size - im.height) // 2)) |
| x = np.asarray(sq, np.float32) / 255 |
| x = (x - np.array(cfg["mean"], np.float32)) / np.array(cfg["std"], np.float32) |
| return torch.from_numpy(x.transpose(2, 0, 1)) |
|
|
|
|
| @torch.inference_mode() |
| def predict(paths, model, cfg): |
| """Logits: above recommended_logit_threshold means a two-page spread.""" |
| return model(torch.stack([preprocess(p, cfg) for p in paths])).tolist() |
|
|
|
|
| if __name__ == "__main__": |
| model, cfg = load_model() |
| thr = cfg["recommended_logit_threshold"] |
| for path, z in zip(sys.argv[1:], predict(sys.argv[1:], model, cfg)): |
| print(f"{path}\t{z:.2f}\t{'spread' if z > thr else 'single'}") |
|
|
