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结构：

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结构：

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