基于自動(dòng)編碼器的賽車視角轉(zhuǎn)換與分割

來(lái)源：Deephub Imba

如下所示，輸入圖像為：

輸出：

總結(jié)來(lái)說(shuō)我們的任務(wù)是獲取輸入圖像，即前方軌道的前置攝像頭視圖，并構(gòu)建一個(gè)鳥(niǎo)瞰軌道視圖，而鳥(niǎo)瞰軌道視圖會(huì)分割不同的顏色表示賽道和路面的邊界。
僅僅從輸入圖像中提取出關(guān)于走向的信息是相當(dāng)困難的，因?yàn)槲磥?lái)的許多軌道信息被壓縮到圖像的前20個(gè)像素行中。鳥(niǎo)瞰攝像頭能夠以更清晰的格式表達(dá)關(guān)于前方賽道的信息，我們可以更容易地使用它來(lái)規(guī)劃汽車的行為。
在正常行駛時(shí)拍攝鳥(niǎo)瞰圖是非常難實(shí)現(xiàn)的，所以如果我們可以使用前置攝像頭的圖像重建這些鳥(niǎo)眼圖像，就能讓我們用更清晰信息來(lái)進(jìn)行路徑的規(guī)劃。另一個(gè)好處是可以降低維度，有效地將整個(gè)圖像表示為一組32個(gè)數(shù)字，這比整個(gè)圖像占用的空間少得多。并且如果還可以使用這種低維數(shù)據(jù)作為強(qiáng)化學(xué)習(xí)算法的觀察空間。
本文中利用一種叫做變分自動(dòng)編碼器(VAEs)的工具來(lái)幫助我們完成這項(xiàng)任務(wù)。簡(jiǎn)單地說(shuō)，我們把圖像壓縮到32維的潛在空間，然后重建我們分割的鳥(niǎo)瞰圖。本文末尾的PyTorch代碼顯示了完整的模型代碼。

為了訓(xùn)練這一點(diǎn)，我們從前置攝像頭和鳥(niǎo)類攝像頭收集了一系列圖像。然后用編碼器進(jìn)行編碼，然后使用全連接的層將維度降低到目標(biāo)大小，最后使用****用一系列反卷積層重建圖像。
結(jié)果如下所示：

雖然我們可以在重建中看到一些噪聲，但它可以很好地捕捉到整體曲線。代碼如下：

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

以上代碼修是從L2R示例代碼進(jìn)行了進(jìn)一步修改，https：//github.com/learn-to-race/l2r

Pytorch中的VAE代碼來(lái)自：

https：//github.com/sksq96/pytorch-vae