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.")