如何调整密度图模型以准确计数视频幼虫（规避点目标误判）

问题背景 #

要完成视频中的幼虫计数任务，最初尝试全帧用密度图算法但误判非幼虫目标，转而先做单帧处理：用K-means聚类确定幼虫颜色生成掩码，再对掩码图像应用密度图算法。但当前预测计数约1000，实际仅约50，模型误将点目标计入密度统计，需优化模型解决该问题。

核心问题分析 #

1. 模型未经过训练：当前代码中的模型是随机初始化的，没有用带幼虫标注的数据集训练，完全是随机预测，必然导致计数严重偏差。
2. 掩码精度不足：仅用alpha通道做掩码，未结合K-means得到的幼虫颜色做精准过滤，仍有大量非幼虫点目标进入模型。
3. 模型结构过于简单：两层Conv无法有效捕捉幼虫的形态特征，难以区分幼虫和噪点。
4. 缺少后处理过滤：未对预测的密度图做噪点过滤，小点目标被直接计入总和。

优化步骤 #

1. 准备训练数据并训练模型 #

密度图模型必须用带标注的数据集训练：

• 对每帧图像标注每个幼虫的中心坐标
• 用高斯核将每个点标注转换为对应的密度图标签（这是密度图计数的核心标签）
• 用标注好的图像-密度图对训练模型，损失函数用MSE（和你当前一致）

2. 优化掩码生成逻辑 #

结合K-means聚类得到的幼虫颜色，生成更精准的颜色掩码：

# 假设K-means得到的幼虫颜色范围是lower_color到upper_color（RGBA）
lower_color = np.array([r_min, g_min, b_min, a_min])
upper_color = np.array([r_max, g_max, b_max, a_max])
# 生成颜色掩码
color_mask = cv2.inRange(image, lower_color, upper_color)
# 结合alpha掩码
final_mask = cv2.bitwise_and(binary_mask, color_mask)

3. 改进模型结构 #

使用更适合密度图预测的轻量结构，加入池化和上采样增强特征提取能力：

def create_density_map_model(input_shape):
    inputs = Input(shape=input_shape)
    # 特征提取
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = tf.keras.layers.MaxPooling2D((2,2))(x)
    
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = tf.keras.layers.MaxPooling2D((2,2))(x)
    
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    
    # 上采样恢复尺寸
    x = tf.keras.layers.UpSampling2D((2,2))(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = tf.keras.layers.UpSampling2D((2,2))(x)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    
    density_map = Conv2D(1, (1, 1), activation='linear', padding='same')(x)
    model = Model(inputs=inputs, outputs=density_map)
    return model

4. 加入后处理过滤 #

对预测的密度图做阈值过滤+连通域分析，去除小噪点：

# 阈值过滤，只保留密度值高于阈值的区域
threshold = 0.01  # 根据实际情况调整
density_thresholded = np.where(predicted_density_map_filtered > threshold, predicted_density_map_filtered, 0)

# 连通域分析，去除面积过小的区域
density_uint8 = (density_thresholded * 255).astype(np.uint8)
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(density_uint8, connectivity=8)

# 过滤小连通域（假设幼虫对应的连通域面积至少为10像素）
min_area = 10
filtered_density = np.zeros_like(density_thresholded)
for i in range(1, num_labels):
    if stats[i, cv2.CC_STAT_AREA] >= min_area:
        filtered_density[labels == i] = density_thresholded[labels == i]

# 计算最终计数
num_larvae_predicted = int(np.sum(filtered_density))

修改后的完整代码 #

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Conv2D, Input, MaxPooling2D, UpSampling2D
from tensorflow.keras.models import Model
from tensorflow.keras.losses import MeanSquaredError
import cv2

# 改进的密度图模型
def create_density_map_model(input_shape):
    inputs = Input(shape=input_shape)
    # 特征提取
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = MaxPooling2D((2,2))(x)
    
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = MaxPooling2D((2,2))(x)
    
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    
    # 上采样恢复尺寸
    x = UpSampling2D((2,2))(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = UpSampling2D((2,2))(x)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    
    density_map = Conv2D(1, (1, 1), activation='linear', padding='same')(x)
    model = Model(inputs=inputs, outputs=density_map)
    return model

# 加载图像
image_path = 'color_17.png'
image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)

# 生成alpha掩码
alpha_channel = image[:, :, 3]
_, binary_mask = cv2.threshold(alpha_channel, 0, 255, cv2.THRESH_BINARY)

# 生成颜色掩码（替换为K-means得到的幼虫颜色范围）
# 示例颜色范围，需根据K-means结果调整
lower_color = np.array([100, 100, 100, 200])
upper_color = np.array([200, 200, 200, 255])
color_mask = cv2.inRange(image, lower_color, upper_color)
# 合并掩码
final_mask = cv2.bitwise_and(binary_mask, color_mask)

image_height, image_width = image.shape[:2]

# 创建模型（注意：这里需要先训练模型，否则还是随机预测）
input_shape = (image_height, image_width, 1)
model = create_density_map_model(input_shape)
model.compile(optimizer='adam', loss=MeanSquaredError())

# 预处理图像
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
image_normalized = image_gray / 255.0
image_input = np.expand_dims(image_normalized, axis=-1)
image_input = np.expand_dims(image_input, axis=0)

# 预测密度图（训练后再运行这部分）
predicted_density_map = model.predict(image_input)
predicted_density_map = np.reshape(predicted_density_map, (image_height, image_width, 1))

# 应用掩码
final_mask_reshaped = np.expand_dims(final_mask, axis=-1)
predicted_density_map_filtered = predicted_density_map * (final_mask_reshaped / 255.0)

# 后处理过滤噪点
threshold = 0.01
density_thresholded = np.where(predicted_density_map_filtered > threshold, predicted_density_map_filtered, 0)

# 连通域分析过滤小区域
density_uint8 = (density_thresholded * 255).astype(np.uint8)
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(density_uint8, connectivity=8)
min_area = 10
filtered_density = np.zeros_like(density_thresholded)
for i in range(1, num_labels):
    if stats[i, cv2.CC_STAT_AREA] >= min_area:
        filtered_density[labels == i] = density_thresholded[labels == i]

# 保存过滤后的密度图
cv2.imwrite('density_plot_filtered.jpg', (filtered_density * 255.0).astype(np.uint8))

# 计算预测数量
num_larvae_predicted = int(np.sum(filtered_density))
print(f"Number of larvae predicted: {num_larvae_predicted}")

辅助参考图像 #

• 未过滤的第一帧：
• 目标幼虫的RGBA颜色示例：
• 当前生成的密度图：

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