(三十一)图形推理与类别名称
# 续写同一文件:training_project/src/readingtorch/predict.py
# 先核对 8×8 像素,再补上批次轴和通道轴;将模型输出转换为概率与类别。
def symbol(self, image):
if self.saved["task"] != "images":
raise ValueError("这份检查点用于页数回归")
image = np.asarray(image, dtype=np.float32)
if image.shape != (8, 8) or not np.isfinite(image).all():
raise ValueError("请提供包含有限数值的 8×8 图像")
x = torch.from_numpy(image.copy()).unsqueeze(0).unsqueeze(0).to(self.device)
with torch.inference_mode():
probabilities = self.model(x).softmax(1)[0].cpu()
index = int(probabilities.argmax())
return {
"label": self.saved["classes"][index],
"probabilities": probabilities.tolist(),
}单张图原为 (8,8),先增加批次和通道两根轴,成为 (1,1,8,8)。模型输出 (1,3),这时为展示结果调用 softmax,取最大概率的编号,再按检查点中的类别表得到“竖线”“横线”或“十字”。
image.copy() 让本次输入独立保存,随后转换为 float32 张量。图像的通道、长宽与有限数值都在这个入口核对,发生错误时直接说明需要的 8×8 输入。