1import torch
2from torchvision import transforms
3from PIL import Image
4
5# 加载模型
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7model = torch.load("best_model.pth", map_location=device)
8model.eval()
9
10# 预测
11transform = transforms.Compose([
12 transforms.Resize((224, 224)),
13 transforms.ToTensor(),
14 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
15])
16
17image = Image.open("your_image.jpg")
18input_tensor = transform(image).unsqueeze(0).to(device)
19
20with torch.no_grad():
21 output = model(input_tensor)
22 probabilities = torch.nn.functional.softmax(output, dim=1)
23 predicted_class = "cat" if probabilities[0][0] > 0.5 else "dog"
24 confidence = max(probabilities[0][0].item(), probabilities[0][1].item()) * 100
25
26print(f"预测结果: {predicted_class}")
27print(f"置信度: {confidence:.2f}%")
1from huggingface_hub import hf_hub_download
2import torch
3
4# 下载模型
5model_path = hf_hub_download(repo_id="wzx952/cat-dog-classifier", filename="best_model.pth")
6model = torch.load(model_path)
7model.eval()
