如何在Google Colab用TensorFlow构建计数红色圆点的神经网络

一、自定义标签的数据集构建（替代image_dataset_from_directory） #

tf.keras.utils.image_dataset_from_directory是为按文件夹分类的标签设计的，没法直接传入你的自定义字典标签。换用tf.data.Dataset手动构建更灵活，步骤如下：

1. 匹配图片路径与标签

   import os
   import tensorflow as tf
   from glob import glob
   
   # 假设图片存储在'./red_dots_imgs'文件夹下
   img_paths = glob('./red_dots_imgs/im_*.png')
   label_dict = {"im_001": 3, "im_002": 5, ...} # 你的标签字典
   
   # 提取文件名并匹配对应标签
   labels = []
   for path in img_paths:
       filename = os.path.basename(path).replace('.png', '')
       labels.append(label_dict[filename])
   

2. 构建并预处理数据集

   def load_preprocess_img(path, label):
       # 加载并解码图片
       img = tf.io.read_file(path)
       img = tf.image.decode_png(img, channels=3)
       # 统一图片尺寸（可根据需求调整，比如224x224）
       img = tf.image.resize(img, (224, 224))
       # 归一化到0-1区间
       img = tf.cast(img, tf.float32) / 255.0
       # 可选：只保留红色通道，过滤其他颜色干扰
       img = img[:, :, 0:1]
       return img, label
   
   # 构建基础数据集
   dataset = tf.data.Dataset.from_tensor_slices((img_paths, labels))
   # 并行加载预处理
   dataset = dataset.map(load_preprocess_img, num_parallel_calls=tf.data.AUTOTUNE)
   
   # 划分训练/验证集，打乱并分批
   train_size = int(0.8 * len(img_paths))
   train_ds = dataset.take(train_size).shuffle(100).batch(32).prefetch(tf.data.AUTOTUNE)
   val_ds = dataset.skip(train_size).batch(32).prefetch(tf.data.AUTOTUNE)
   

二、计数模型的层组合（无需层数等于最大圆点数量） #

圆点计数是回归任务（预测连续数值），不是分类任务，完全不需要让层数等于最大圆点数量。推荐用「CNN特征提取+全连接层回归」的结构：

示例模型结构 #

model = tf.keras.Sequential([
    # 特征提取：卷积+池化堆叠
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(224,224,1)), # 单通道输入对应(224,224,1)
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2,2)),
    # 扁平化特征图为一维向量
    tf.keras.layers.Flatten(),
    # 全连接层做特征映射
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(64, activation='relu'),
    # 输出层：回归任务用单神经元，默认激活函数为linear
    tf.keras.layers.Dense(1)
])

# 编译模型：回归任务用均方误差损失
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

关键说明 #

• 卷积层负责提取圆点的形状、颜色特征，堆叠层数根据图片复杂度调整（3-5层足够应对大多数场景）
• 输出层仅需1个神经元，直接预测圆点数量，无需对应每个可能的计数类别
• 损失函数选mse（均方误差），评估指标用mae（平均绝对误差）更直观，代表预测值与真实值的平均偏差

额外优化建议 #

• 数据增强：如果圆点的位置、大小存在差异，可添加tf.keras.layers.RandomFlip、RandomRotation等增强层，提升模型泛化能力
• 红色通道预处理：只保留红色通道能过滤其他颜色干扰，让模型更专注于红色圆点的特征

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