Spaces:
Sleeping
Sleeping
| import json | |
| from pathlib import Path | |
| import torch | |
| from safetensors.torch import load_file | |
| from torchvision import transforms | |
| from torchvision.models import ResNet18_Weights, resnet18 | |
| MEAN = [0.485, 0.456, 0.406] | |
| STD = [0.229, 0.224, 0.225] | |
| def build_model(num_classes: int, pretrained: bool = False): | |
| model = resnet18(weights=ResNet18_Weights.DEFAULT if pretrained else None) | |
| model.fc = torch.nn.Linear(model.fc.in_features, num_classes) | |
| return model | |
| def preprocessing(config: dict, training: bool = False): | |
| size = config["image_size"] | |
| operations = ( | |
| [transforms.RandomResizedCrop(size, scale=(.8, 1.0)), transforms.RandomHorizontalFlip()] | |
| if training else | |
| [transforms.Resize(round(size * 256 / 224)), transforms.CenterCrop(size)] | |
| ) | |
| return transforms.Compose(operations + [ | |
| transforms.ToTensor(), | |
| transforms.Normalize(config["mean"], config["std"]), | |
| ]) | |
| def load_classifier(directory: str | Path): | |
| directory = Path(directory) | |
| config = json.loads((directory / "config.json").read_text()) | |
| model = build_model(len(config["labels"])) | |
| model.load_state_dict(load_file(str(directory / "model.safetensors"))) | |
| model.eval() | |
| return model, config | |
| def classify(model, config: dict, image, top_k: int): | |
| if image is None: | |
| raise ValueError("Choose an image") | |
| if not 1 <= top_k <= len(config["labels"]): | |
| raise ValueError("Invalid top_k") | |
| tensor = preprocessing(config)(image.convert("RGB")).unsqueeze(0) | |
| tensor = tensor.to(next(model.parameters()).device) | |
| probabilities = model(tensor).softmax(dim=1)[0] | |
| scores, indices = probabilities.topk(top_k) | |
| return [ | |
| {"label": config["labels"][index], "score": float(score)} | |
| for score, index in zip(scores.tolist(), indices.tolist()) | |
| ] | |