基于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 在评论区打出来,我会去尝试~

相关推荐
weixin_515202495 小时前
第R4周-LSTM-火灾温度预测
人工智能·rnn·lstm
沅_Yuan1 天前
基于GRU门控循环神经网络的多分类预测【MATLAB】
matlab·分类·gru
数据分析能量站2 天前
目标检测-R-CNN
目标检测·r语言·cnn
一只Black2 天前
LSTM实现天气模型训练与预测
人工智能·rnn·lstm
落魄君子2 天前
CNN回归-卷积神经网络(Convolutional Neural Network)
人工智能·神经网络·回归·cnn
机器学习之心3 天前
BiTCN-BiGRU基于双向时间卷积网络结合双向门控循环单元的数据多特征分类预测(多输入单输出)
深度学习·分类·gru
机器学习之心3 天前
Bayes-GRU-Attention的数据多特征分类预测Matlab实现
matlab·分类·gru
机器学习之心4 天前
LSTM-SVM时序预测 | Matlab基于LSTM-SVM基于长短期记忆神经网络-支持向量机时间序列预测
神经网络·支持向量机·lstm
love you joyfully4 天前
目标检测与R-CNN——pytorch与paddle实现目标检测与R-CNN
人工智能·pytorch·目标检测·cnn·paddle
Debroon4 天前
乳腺癌多模态诊断解释框架:CNN + 可解释 AI 可视化
人工智能·神经网络·cnn