torch_pretrain_legacy.py

import torch
import torchvision.models as models
import torchvision.transforms as transforms

from PIL import Image


img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

# Step 1: Load a pre-trained model.
# In this step we will load a ResNet architecture.
model = models.resnet50(pretrained=True)
# Set the model in evaluation mode.
model.eval()

# Step 2: Define and initialize a composition
# of data transformations.
preprocess = transforms.Compose([
    transforms.Resize([256, ]),
    transforms.CenterCrop(224),
    transforms.PILToTensor(),
    transforms.ConvertImageDtype(torch.float),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Step 3: Get the predictions on the (processed)
# test dataset.
batch = preprocess(img).unsqueeze(0)
prediction = model(batch).squeeze(0).softmax(0)

# Step 4: Print a human-readable output
class_id = prediction.argmax().item()
score = prediction[class_id].item()
with open("imagenet_classes.txt", "r") as f:
    categories = [s.strip() for s in f.readlines()]
    category_name = categories[class_id]
print(f"{category_name}: {100 * score}%")