06_image_classification/main.py

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()