基于CNN+RNNs(LSTM, GRU)的红点位置检测(pytorch)

1 项目背景

需要在图片精确识别三跟红线所在的位置,并输出这三个像素的位置。

其中,每跟红线占据不止一个像素,并且像素颜色也并不是饱和度和亮度极高的红黑配色,每个红线放大后可能是这样的。

而我们的目标是精确输出每个红点的位置,需要精确到像素。也就是说,对于每根红线,模型需要输出橙色箭头所指的像素而不是蓝色箭头所指的像素的位置。

在之前尝试过纯 RNNs 检测红点,但是准确率感人,在噪声极低的情况下并不能精准识别位置。但是有次尝试transformer位置编码之后发现效果不错:

实验 loss 完全准确的点
GRU 129.6641 1762.0/9000 (20%)
LSTM 249.2053 1267.0/9000 (14%)
Position embedding + GRU 16.3403 5025.0/9000 (56%)
Position embedding + LSTM 204.1551 1603.0/9000 (18%)

这说明模型的难点在于学习位置信息而不是寻找颜色有问题的点。联想到CNN也能提供位置信息,我决定尝试卷积一下的效果。

2 数据集

还是之前那个代码合成的数据集数据集,每个数据集规模在15000张图片左右,在没有加入噪音的情况下,每个样本预览如图所示:

加入噪音后,每个样本的预览如下图所示:

图中黑色部分包含比较弱的噪声,并非完全为黑色。

数据集包含两个文件,一个是文件夹,里面包含了jpg压缩的图像数据:

另一个是csv文件,里面包含了每个图像的名字以及3根红线所在的像素的位置。

3 思路

其实思路特别朴素。就是在RNNs要读序列化数据之前先用CNN把数据跑一遍,让原始的输入序列变成具有局部特征表示的嵌入表示,卷积后提取的特征输入到 RNN层,RNN 保持了序列中的长时依赖信息。接下来先用 fc1 把 RNN 的输出映射成分数,然后用 fc2 预测三个具体位置,经过 Sigmoid 输出 [0, 1] 的相对位置,再与宽度相乘得到真实位置。具体的流程如下图所示:

4 结果

在图片长度为1080、低噪声环境时,对比实验的结果如下:

实验 loss 完全准确的点
GRU 129.6641 1762.0/9000 (20%)
LSTM 249.2053 1267.0/9000 (14%)
CNN+GRU 1419.5781 601.0/9000 (7%)
CNN+LSTM 1166.4599 762.0/9000 (8%)

1080长度下图片抽样预测的效果如下:

在简单图片中的效果跟其他方法差距不大------基本都能准确定位红线,但是还是没办法做到像素级别的精确

可能是我的打开方式不对,但是CNN+RNN的效果并不如意。

从训练过程来看存在过拟合:

5 代码

CNN+GRU结构:

python 复制代码
class CNN_GRU(nn.Module):
    def __init__(self, config):
        super(CNN_GRU, self).__init__()
        self.input_size = config.input_size
        self.hidden_size = config.hidden_size
        self.num_layers = config.num_layers
        self.device = config.device

        # CNN
        self.conv1 = nn.Conv1d(in_channels=self.input_size, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(in_channels=128, out_channels=self.input_size, kernel_size=3, padding=1)

        self.gru = nn.GRU(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=self.num_layers,
                            batch_first=True, bidirectional=True, dropout=0.6)

        self.fc1 = nn.Sequential(
            nn.Linear(self.hidden_size * 2, 1)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(config.width, 3),  # predict 3 points
            nn.Sigmoid(),
        )

        self.scale = config.width
        self.device = config.device

    def forward(self, x):
        x = x.squeeze(2)
        x = F.relu(self.conv1(x))  # (batch_size, 64, width)
        x = F.relu(self.conv2(x))  # (batch_size, 128, width)
        x = F.relu(self.conv3(x))  # (batch_size, input_size, width)
      	
      	x = x.permute(0, 2, 1)
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
        output, _ = self.gru(x0, h0)
        
        scores = self.fc1(output).squeeze(-1)  # shape: (batch_size, 1080)
        predicted_positions = self.fc2(scores)
        
        scaled_predicted_positions = predicted_positions * self.scale
        final_predicted_positions = torch.clamp(scaled_predicted_positions, min=0, max=self.scale - 1)
        return final_predicted_positions

CNN+LSTM结构:

python 复制代码
class CNN_GRU(nn.Module):
    def __init__(self, config):
        super(CNN_GRU, self).__init__()
        self.input_size = config.input_size
        self.hidden_size = config.hidden_size
        self.num_layers = config.num_layers
        self.device = config.device
        # CNN
        self.conv1 = nn.Conv1d(in_channels=self.input_size, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(in_channels=128, out_channels=self.input_size, kernel_size=3, padding=1)
        self.lstm = nn.LSTM(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=self.num_layers,
                            batch_first=True, bidirectional=True, dropout=0.6)
        self.fc1 = nn.Sequential(
            nn.Linear(self.hidden_size * 2, 1)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(config.width, 3),  # predict 3 points
            nn.Sigmoid(),
        )
        self.scale = config.width
        self.device = config.device

    def forward(self, x):
        x = x.squeeze(2)
        x = F.relu(self.conv1(x))  # (batch_size, 64, width)
        x = F.relu(self.conv2(x))  # (batch_size, 128, width)
        x = F.relu(self.conv3(x))  # (batch_size, input_size, width)
      
      	x = x.permute(0, 2, 1)
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
        output, _ = self.lstm(x, (h0, c0))
        
        scores = self.fc1(output).squeeze(-1)  # shape: (batch_size, 1080)
        predicted_positions = self.fc2(scores)
        
        scaled_predicted_positions = predicted_positions * self.scale
        final_predicted_positions = torch.clamp(scaled_predicted_positions, min=0, max=self.scale - 1)
        return final_predicted_positions

路过的大佬有什么建议 ball ball 在评论区打出来,我会去尝试~

相关推荐
Evaporator Core2 小时前
门控循环单元(GRU)与时间序列预测应用
人工智能·深度学习·gru
Matlab精灵3 小时前
基于CNN-LSTM的时序预测MATLAB实战
matlab·cnn·lstm
爱研究的小牛4 小时前
AIVA 技术浅析(四):捕捉音乐作品中的长期依赖关系
人工智能·rnn·深度学习·aigc·lstm
AI浩4 小时前
ShuffleNet:一种为移动设备设计的极致高效的卷积神经网络
人工智能·神经网络·cnn
卧式纯绿11 小时前
自动驾驶3D目标检测综述(三)
人工智能·python·深度学习·目标检测·3d·cnn·自动驾驶
白光白光21 小时前
量子卷积神经网络
人工智能·神经网络·cnn
苏涵.1 天前
深度学习实验十二 卷积神经网络(3)——基于残差网络实现手写体数字识别实验
人工智能·深度学习·神经网络·cnn
机器学习之心2 天前
一区北方苍鹰算法优化+创新改进Transformer!NGO-Transformer-LSTM多变量回归预测
算法·lstm·transformer·北方苍鹰算法优化·多变量回归预测·ngo-transformer
阿_旭3 天前
TensorFlow构建CNN卷积神经网络模型的基本步骤:数据处理、模型构建、模型训练
人工智能·深度学习·cnn·tensorflow