基于Timm的图像预测Docker镜像开发求助：JSON输出与403错误

1. JSON输出无效问题 #

你的代码直接打印Python字典，输出的是单引号格式，不符合JSON要求（必须使用双引号）。要生成标准JSON，需用json.dumps()序列化后再打印：

修改后的输出代码片段： #

# 获取模型对应的真实类别标签（替换硬编码的"cat"）
class_names = model.pretrained_cfg['class_names']
top_prob, top_idx = torch.topk(probabilities, 1)
predicted_class = class_names[top_idx.item()]
confidence = f"{top_prob.item():.2f}"

# 生成标准JSON并打印
result = {"predicted": predicted_class, "confidence": confidence}
print(json.dumps(result))

2. urllib 403 Forbidden错误 #

这是因为目标服务器拒绝了无标识的爬虫请求，需要给请求添加User-Agent头模拟浏览器访问。替换原有的urllib.request.urlretrieve代码：

修改后的图片下载代码片段： #

url = args.image
filename = "temp_image.jpg"

# 添加请求头，模拟Chrome浏览器
headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
}
req = urllib.request.Request(url, headers=headers)

with urllib.request.urlopen(req) as response, open(filename, 'wb') as out_file:
    out_file.write(response.read())

img = img.open(filename).convert('RGB')

完整修正后的代码 #

from __future__ import print_function
import argparse
import torch
import timm
import urllib
import json
from PIL import Image as img
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

# Training settings
parser = argparse.ArgumentParser(description='Assignment 1')
parser.add_argument('--model', type=str)
parser.add_argument('--image', type=str)

if __name__ == '__main__':
    args = parser.parse_args()
    model = timm.create_model(args.model, pretrained=True)
    model.eval()
    config = resolve_data_config({}, model=model)
    transform = create_transform(**config)

    # 修正图片下载逻辑，添加User-Agent
    url = args.image
    filename = "temp_image.jpg"
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }
    req = urllib.request.Request(url, headers=headers)
    with urllib.request.urlopen(req) as response, open(filename, 'wb') as out_file:
        out_file.write(response.read())

    img = img.open(filename).convert('RGB')
    tensor = transform(img).unsqueeze(0)  # transform and add batch dimension

    with torch.no_grad():
        out = model(tensor)
    probabilities = torch.nn.functional.softmax(out[0], dim=0)

    # 获取真实预测类别，生成标准JSON
    class_names = model.pretrained_cfg['class_names']
    top_prob, top_idx = torch.topk(probabilities, 1)
    predicted_class = class_names[top_idx.item()]
    confidence = f"{top_prob.item():.2f}"

    result = {"predicted": predicted_class, "confidence": confidence}
    print(json.dumps(result))

额外说明 #

• 原代码中硬编码"cat"作为预测结果，修正后会根据模型输出自动匹配对应类别
• 使用:.2f格式化置信度，确保输出两位小数，和示例格式一致

内容的提问来源于stack exchange，提问作者3PO2