File size: 1,232 Bytes
93baf52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import argparse, json, torch
from torchvision import models, transforms
from PIL import Image
import urllib.request
IMAGENET_URL = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
def load_labels():
with urllib.request.urlopen(IMAGENET_URL) as f:
labels = [s.strip() for s in f.read().decode("utf-8").splitlines()]
return labels
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--image", type=str, default=None, help="Path to an image")
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT).to(device).eval()
preprocess = models.MobileNet_V2_Weights.DEFAULT.transforms()
img = Image.open(args.image).convert("RGB") if args.image else Image.new("RGB", (224,224), "white")
x = preprocess(img).unsqueeze(0).to(device)
with torch.no_grad():
logits = model(x)
probs = torch.softmax(logits, dim=-1)[0]
top5 = torch.topk(probs, 5)
labels = load_labels()
for p, idx in zip(top5.values, top5.indices):
print(f"{labels[idx]}: {float(p):.4f}")
if __name__ == "__main__":
main()
|