基于自动编码器的赛车视角转换与分割

本篇文章将介绍如何将赛道的图像转换为语义分割后鸟瞰图的轨迹。

如下所示，输入图像为

输出：

总结来说我们的任务是获取输入图像，即前方轨道的前置摄像头视图，并构建一个鸟瞰轨道视图，而鸟瞰轨道视图会分割不同的颜色表示赛道和路面的边界。

仅仅从输入图像中提取出关于走向的信息是相当困难的，因为未来的许多轨道信息被压缩到图像的前20个像素行中。鸟瞰摄像头能够以更清晰的格式表达关于前方赛道的信息，我们可以更容易地使用它来规划汽车的行为。

在正常行驶时拍摄鸟瞰图是非常难实现的，所以如果我们可以使用前置摄像头的图像重建这些鸟眼图像，就能让我们用更清晰信息来进行路径的规划。另一个好处是可以降低维度，有效地将整个图像表示为一组32个数字，这比整个图像占用的空间少得多。并且如果还可以使用这种低维数据作为强化学习算法的观察空间。

本文中利用一种叫做变分自动编码器(VAEs)的工具来帮助我们完成这项任务。简单地说，我们把图像压缩到32维的潜在空间，然后重建我们分割的鸟瞰图。本文末尾的PyTorch代码显示了完整的模型代码。

为了训练这一点，我们从前置摄像头和鸟类摄像头收集了一系列图像。然后用编码器进行编码，然后使用全连接的层将维度降低到目标大小，最后使用解码器用一系列反卷积层重建图像。

结果如下所示：

虽然我们可以在重建中看到一些噪声，但它可以很好地捕捉到整体曲线。代码如下：

import cv2
import tqdm
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class BEVVAE(nn.Module):
    """Input should be (bsz, C, H, W) where C=3, H=42, W=144"""

    def __init__(self, im_c=3, im_h=95, im_w=512, z_dim=32):
        super().__init__()

        self.im_c = im_c
        self.im_h = im_h
        self.im_w = im_w

        encoder_list = [
            nn.Conv2d(im_c, 32, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Flatten(),
        ]
        self.encoder = nn.Sequential(*encoder_list)
        self.encoder_list = encoder_list
        sample_img = torch.zeros([1, im_c, im_h, im_w])
        em_shape = nn.Sequential(*encoder_list[:-1])(sample_img).shape[1:]
        h_dim = np.prod(em_shape)

        self.fc1 = nn.Linear(h_dim, z_dim)
        self.fc2 = nn.Linear(h_dim, z_dim)
        self.fc3 = nn.Linear(z_dim, h_dim)

        self.decoder = nn.Sequential(
            nn.Unflatten(1, em_shape),
            nn.ConvTranspose2d(
                em_shape[0],
                256,
                kernel_size=4,
                stride=2,
                padding=1,
                output_padding=(1, 0),
            ),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, output_padding=(1, 0)),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, output_padding=(1, 0)),
            nn.ReLU(),
            nn.ConvTranspose2d(
                64, 32, kernel_size=4, stride=2, padding=1, output_padding=(1, 0)
            ),
            nn.ReLU(),
            nn.ConvTranspose2d(32, im_c, kernel_size=4, stride=2, padding=1, output_padding=(1, 0)),
            nn.Sigmoid(),
        )

    def reparameterize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        esp = torch.randn(*mu.size(), device=mu.device)
        z = mu + std * esp
        return z

    def bottleneck(self, h):
        mu, logvar = self.fc1(h), self.fc2(h)
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def representation(self, x):
        return self.bottleneck(self.encoder(x))[0]

    def encode_raw(self, x: np.ndarray, device):
        # assume x is RGB image with shape (bsz, H, W, 3)
        p = np.zeros([x.shape[0], 95, 512, 3], np.float)
        for i in range(x.shape[0]):
            p[i] = x[i][190:285] / 255
        x = p.transpose(0, 3, 1, 2)
        x = torch.as_tensor(x, device=device, dtype=torch.float)
        v = self.representation(x)
        return v, v.detach().cpu().numpy()

    def squish_targets(self, x: np.ndarray) -> np.ndarray:
        # Take in target images and resize them
        p = np.zeros([x.shape[0], 95, 512, 3], np.float)
        for i in range(x.shape[0]):
            p[i] = cv2.resize(x[i], (512, 95)) / 255
        x = p.transpose(0, 3, 1, 2)
        return x

    def encode(self, x):
        h = self.encoder(x)
        z, mu, logvar = self.bottleneck(h)
        return z, mu, logvar

    def decode(self, z):
        z = self.fc3(z)
        return self.decoder(z)

    def forward(self, x):
        # expects (N, C, H, W)
        z, mu, logvar = self.encode(x)
        z = self.decode(z)
        return z, mu, logvar

    def loss(self, bev, recon, mu, logvar, kld_weight=1.0):
        bce = F.binary_cross_entropy(recon, bev, reduction="sum")
        kld = -0.5 * torch.sum(1 + logvar - mu ** 2 - logvar.exp())
        return bce + kld * kld_weight

https://avoid.overfit.cn/post/48f129f8e05242128cc55be13433ad0a

作者：Nandan Tumu