TensorFlow模型检测特定类别时执行自定义操作的问题

问题概述 #

开发了一款可检测人脸朝向（左侧、中间、右侧）的TensorFlow模型，预测代码能正常输出类别及置信度，但添加自定义逻辑时出现异常：当检测到左侧/右侧人脸时，无法触发print('face the computer')的操作，总是进入else分支输出无效提示，即使移除else，if语句也会被跳过。

原因分析 #

问题出在标签读取时的换行符：使用readlines()读取labels.txt时，每一行末尾会保留换行符\n。例如检测到右侧人脸时，class_name的实际值是"2 Right_Side_Face\n"，而判断条件里写的是"2 Right_Side_Face"，两者不相等，导致条件不成立。从测试输出里也能看出：Class: 2 Right_Side_Face之后直接换行了，说明class_name自带换行符。

修复方案 #

提供两种可靠的修复方式：

方式1：读取标签时清除换行符 #

修改标签读取代码，对每一行执行strip()操作，去除首尾的换行符和空白字符：

# Load the labels
class_names = [line.strip() for line in open('models//labels.txt', 'r').readlines()]

之后原有的字符串判断逻辑就能正常工作：

if class_name == "2 Right_Side_Face" or class_name == "0 Left_Side_Face":
    print('face the computer')
else:
    print('Invalid class name')
    print('exiting.....')

方式2：直接使用预测索引判断（更高效） #

既然已经通过np.argmax(prediction)拿到了类别索引index，直接用索引判断更可靠，还能避免字符串匹配的潜在问题：

# 0对应左侧人脸，2对应右侧人脸
if index in (0, 2):
    print('face the computer')
else:
    print('Invalid class name')
    print('exiting.....')

完整修复代码示例 #

以下是采用方式2的完整代码：

from keras.models import load_model
from PIL import Image, ImageOps
import numpy as np

# Disable scientific notation for clarity
np.set_printoptions(suppress=True)

# Load the model
model = load_model('models//keras_Model.h5', compile=False)

# Load the labels
class_names = open('models//labels.txt', 'r').readlines()

# Create the array of the right shape to feed into the keras model
data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)

# Replace this with the path to your image
image = Image.open('IMG.png').convert('RGB')

# resize the image to a 224x224
size = (224, 224)
image = ImageOps.fit(image, size, Image.Resampling.LANCZOS)

# turn the image into a numpy array
image_array = np.asarray(image)

# Normalize the image
normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1

# Load the image into the array
data[0] = normalized_image_array

# run the inference
prediction = model.predict(data)
index = np.argmax(prediction)
class_name = class_names[index]
confidence_score = prediction[0][index]

print('Class:', class_name, end='')
print('Confidence score:', confidence_score)

# 修复后的判断逻辑
if index in (0, 2):
    print('face the computer')
else:
    print('Invalid class name')
    print('exiting.....')

测试验证 #

运行修复后的代码，当检测到左侧或右侧人脸时，会正常打印face the computer；检测到中间人脸时，进入else分支输出对应提示。

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