基于Python的深度学习二分类代码：识别图像是否为苹果

这个方案不需要你做复杂的多类别标注，只要把现有数据集按文件夹归成两类就能直接跑，CPU也能训练出可用的效果。

前置依赖安装 #

先安装需要的第三方库，命令行执行：
pip install tensorflow pillow scikit-learn

数据集整理要求 #

你不需要重新给图片打标签，只要把现有图片按下面的目录结构放好就行：

fruit_dataset/
├─ apple/                # 所有苹果类的图片全部存入这个文件夹
│  ├─ apple_001.jpg
│  ├─ apple_002.jpg
│  └─ ...
└─ non_apple/            # 西红柿、西瓜、橙子等所有非苹果类图片全存入这个文件夹
   ├─ tomato_001.jpg
   ├─ watermelon_001.jpg
   └─ ...

框架会自动把两个文件夹的图片分别标注为「苹果」「非苹果」，不需要手动写标注文件。

完整可运行代码 #

import tensorflow as tf
from tensorflow.keras import layers, models

# -------------------------- 配置参数 --------------------------
DATASET_PATH = "./fruit_dataset"  # 修改为你本地的数据集实际路径
IMAGE_SIZE = (160, 160)
BATCH_SIZE = 16
EPOCHS = 10
CLASS_NAMES = ["非苹果", "苹果"]

# -------------------------- 加载数据集 --------------------------
# 自动按9:1比例拆分训练集、验证集，无需手动拆分
train_ds = tf.keras.utils.image_dataset_from_directory(
    DATASET_PATH,
    validation_split=0.1,
    subset="training",
    seed=123,
    image_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    label_mode="binary"  # 二分类任务专用标注模式
)

val_ds = tf.keras.utils.image_dataset_from_directory(
    DATASET_PATH,
    validation_split=0.1,
    subset="validation",
    seed=123,
    image_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    label_mode="binary"
)

# 数据预处理+增强，降低过拟合概率
data_augmentation = models.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
    layers.Rescaling(1./255)  # 图像像素值归一化到0-1区间
])

# -------------------------- 构建模型 --------------------------
# 采用轻量预训练模型做迁移学习，CPU环境也能快速完成训练
base_model = tf.keras.applications.MobileNetV2(
    input_shape=IMAGE_SIZE + (3,),
    include_top=False,
    weights="imagenet"
)
base_model.trainable = False  # 冻结预训练层权重，仅训练上层分类头

model = models.Sequential([
    data_augmentation,
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.Dropout(0.2),
    layers.Dense(1, activation="sigmoid")  # 二分类任务输出层
])

model.compile(
    optimizer="adam",
    loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=["accuracy"]
)

# -------------------------- 启动训练 --------------------------
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=EPOCHS
)

# -------------------------- 单图预测示例 --------------------------
def predict_image(img_path):
    img = tf.keras.utils.load_img(img_path, target_size=IMAGE_SIZE)
    img_array = tf.keras.utils.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)
    
    pred_score = model.predict(img_array)[0][0]
    print(f"预测结果：{CLASS_NAMES[1 if pred_score > 0.5 else 0]}，置信度：{max(pred_score, 1-pred_score):.2f}")

# 调用示例，把路径换成你要测试的图片路径即可
# predict_image("./test_apple.jpg")

调优提示 #

• 如果你的电脑配置较低，把BATCH_SIZE调整为8即可，训练速度会有所下降，但不会出现内存占满的问题
• 如果验证集准确率低于85%，可以把EPOCHS调整为15_{20，或者给两个类别各补充100}200张不同拍摄角度、不同光线条件的图片
• 训练完成后可以执行model.save("apple_classifier.h5")保存模型，下次使用时直接用tf.keras.models.load_model("apple_classifier.h5")加载，无需重新训练

内容的提问来源于stack exchange，提问作者김수환