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

相关推荐
爱吃泡芙的小白白7 小时前
突破传统:CNN卷积层(普通/空洞)核心技术演进与实战指南
人工智能·神经网络·cnn·卷积层·空洞卷积·普通卷积
机器学习之心HML12 小时前
多光伏电站功率预测新思路:当GCN遇见LSTM,解锁时空预测密码,python代码
人工智能·python·lstm
yuanyuan2o213 小时前
【深度学习】全连接、卷积神经网络
人工智能·深度学习·cnn
Liue6123123115 小时前
自卸车多部件识别 _ Mask R-CNN改进模型实现(Caffe+FPN)_1
r语言·cnn·caffe
爱吃泡芙的小白白18 小时前
深入解析CNN中的BN层:从稳定训练到前沿演进
人工智能·神经网络·cnn·梯度爆炸·bn·稳定模型
水月wwww1 天前
【深度学习】卷积神经网络
人工智能·深度学习·cnn·卷积神经网络
摘星编程1 天前
CANN ops-nn Pooling算子解读:CNN模型下采样与特征提取的核心
人工智能·神经网络·cnn
慢半拍iii1 天前
从零搭建CNN:如何高效调用ops-nn算子库
人工智能·神经网络·ai·cnn·cann
B站_计算机毕业设计之家2 天前
豆瓣电影数据采集分析推荐系统 | Python Vue Flask框架 LSTM Echarts多技术融合开发 毕业设计源码 计算机
vue.js·python·机器学习·flask·echarts·lstm·推荐算法
偷吃的耗子2 天前
【CNN算法理解】:CNN平移不变性详解:数学原理与实例
人工智能·算法·cnn