| """Standalone inference example for adeIaide/skin-lesion-7class-efficientnet.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
| from safetensors.torch import load_file |
| from torch import nn |
| from torchvision import models, transforms |
|
|
|
|
| REPO_ID = "adeIaide/skin-lesion-7class-efficientnet" |
|
|
|
|
| class EfficientNetB0(nn.Module): |
| def __init__(self, num_classes: int = 7, dropout: float = 0.3) -> None: |
| super().__init__() |
| network = models.efficientnet_b0(weights=None) |
| embedding_dim = network.classifier[1].in_features |
| self.encoder = nn.Sequential(network.features, network.avgpool, nn.Flatten(1)) |
| self.dropout = nn.Dropout(dropout) |
| self.classifier = nn.Linear(embedding_dim, num_classes) |
|
|
| def forward(self, image: torch.Tensor) -> torch.Tensor: |
| return self.classifier(self.dropout(self.encoder(image))) |
|
|
|
|
| def load_model(device: torch.device) -> tuple[nn.Module, dict]: |
| config_path = hf_hub_download(REPO_ID, "config.json") |
| weights_path = hf_hub_download(REPO_ID, "model.safetensors") |
| config = json.loads(Path(config_path).read_text(encoding="utf-8")) |
| model = EfficientNetB0( |
| num_classes=int(config["num_classes"]), |
| dropout=float(config["dropout"]), |
| ) |
| model.load_state_dict(load_file(weights_path, device=str(device))) |
| return model.to(device).eval(), config |
|
|
|
|
| def predict(image_path: str | Path) -> dict[str, float]: |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model, config = load_model(device) |
| norm = config["normalization"] |
| transform = transforms.Compose( |
| [ |
| transforms.Resize((config["image_size"], config["image_size"])), |
| transforms.ToTensor(), |
| transforms.Normalize(norm["mean"], norm["std"]), |
| ] |
| ) |
| image = Image.open(image_path).convert("RGB") |
| batch = transform(image).unsqueeze(0).to(device) |
| with torch.inference_mode(): |
| probabilities = torch.softmax(model(batch), dim=1)[0].cpu().tolist() |
| return dict(zip(config["class_names"], probabilities)) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("image", help="Path to one dermoscopic image") |
| args = parser.parse_args() |
| probabilities = predict(args.image) |
| predicted = max(probabilities, key=probabilities.get) |
| print(f"prediction: {predicted}") |
| for label, probability in sorted( |
| probabilities.items(), key=lambda item: item[1], reverse=True |
| ): |
| print(f"{label:5s} {probability:.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
