Create infer.py

infer.py

@@ -0,0 +1,49 @@
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from transformers import AutoModelForImageClassification
+import requests
+from io import BytesIO
+
+MEAN = [0.5, 0.5, 0.5]
+STD  = [0.5, 0.5, 0.5]
+
+transform = T.Compose([
+    T.Resize((32, 32), interpolation=T.InterpolationMode.BICUBIC),
+    T.ToTensor(),
+    T.Normalize(mean=MEAN, std=STD)
+])
+
+def load_image(source: str) -> Image.Image:
+    if source.startswith("http://") or source.startswith("https://"):
+        response = requests.get(source)
+        return Image.open(BytesIO(response.content)).convert("RGB")
+    return Image.open(source).convert("RGB")
+
+def predict(model, image_source: str, device: str = "cpu") -> dict:
+    image = load_image(image_source)
+    x = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = model(pixel_values=x).logits
+    probs = torch.softmax(logits, dim=-1)[0]
+    top5 = probs.topk(5)
+    return {
+        model.config.id2label[i.item()]: f"{p.item()*100:.2f}%"
+        for i, p in zip(top5.indices, top5.values)
+    }
+
+if __name__ == "__main__":
+    import sys
+    source = sys.argv[1] if len(sys.argv) > 1 else "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"  # an ant for fallback!
+    
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    model = AutoModelForImageClassification.from_pretrained(
+        "MamaPearl/nula-cifar10-robust-v0",
+        trust_remote_code=True
+    ).to(DEVICE)
+    model.eval()
+    
+    results = predict(model, source, device=DEVICE) 
+    for label, prob in results.items():
+        print(f"{label:15} {prob}")