File size: 1,697 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import torch
from torchvision import models
from PIL import Image
import urllib.request
import os
# URL for ImageNet class labels
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
# Device selection
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load model
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT).to(device).eval()
preprocess = models.MobileNet_V2_Weights.DEFAULT.transforms()
# Online image
online_image_url = "https://upload.wikimedia.org/wikipedia/commons/9/9a/Pug_600.jpg"
online_image_path = "online_image.jpg"
urllib.request.urlretrieve(online_image_url, online_image_path)
# Offline image from same directory
offline_image_path = "remiai.png" # Replace with your actual image filename
# Function to run inference
def classify_image(image_path):
img = Image.open(image_path).convert("RGB")
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()
print(f"Results for: {image_path}")
for p, idx in zip(top5.values, top5.indices):
print(f"{labels[idx]}: {float(p):.4f}")
print()
# Run inference on both images
classify_image(online_image_path)
if os.path.exists(offline_image_path):
classify_image(offline_image_path)
else:
print(f"Offline image '{offline_image_path}' not found.")
|