【深度學習】用CNN實現(xiàn)全景圖像語義分割！

相信許多讀者體驗過b站上的全景視頻，如果還沒有，快來體驗一下吧[1]！只需鼠標點擊并移動，便可360度無死角的瀏覽全景視頻，讓人如同身臨其境。全景圖像，又稱360°全景圖，其數(shù)據(jù)分布在球面空間上。但是，當我們將全景圖像展開時，會造成畸變。

怎么處理？直接將傳統(tǒng)二維平面圖像處理方法應用到球面數(shù)據(jù)上，其效果則會大大降低。而要解決分布在球面空間上的數(shù)據(jù)，需要特定的方法，比如球面卷積網(wǎng)絡。本文手把手帶你實踐一個有趣的應用——全景圖像語義分割，使用多種傳統(tǒng)CNN方法和球面CNN方法進行對比。

如下圖所示，全景圖分割實例像素級別分類，每種實例對應一個標簽。完成本教程后，你將能夠做一個圖中所示的全景圖小應用。

文章數(shù)據(jù)集獲取與代碼地址見文末

1. 環(huán)境構建

基于深度學習的編程環(huán)境往往有各種復雜的環(huán)境依賴，而各種安裝報錯總是消磨我們的時間，其實之一過程可以大大縮短。我們所需要的也就是通過一個命令安裝所有的依賴并打開環(huán)境

make up #等價于 docker-compose up -d

再通過一個命令

make in

來進入我們需要的環(huán)境，然后運行程序。為實現(xiàn)構建這一過程，基于docker –docker-compose – make來搭建我們的環(huán)境，其原理如下圖所示：

docker –docker-compose – make三個工具對應三個配置文件，都在項目根目錄進行了聲明：

Dockerfile
docker-compose.yml
Makefile

其中

• Dockerfile 定義了實驗所需要的所有環(huán)境，依據(jù)此文件可以編譯成docker鏡像，其中包含我們需要的庫
• docker-compose.yml定義了鏡像的啟動方式，在本文中，我們定義兩個服務，一個作為終端來運行命令，一個作為jupyter lab供調(diào)試
• Makefile定義了啟動環(huán)境的方式

本文實驗環(huán)境：Ubuntu20.04，CUDA11.0，Pytorch1.7

Docker安裝

# 1.安裝docker
sudo apt install -y docker docker.io
# 2.安裝英偉達docker
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
&& curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
# 3.安裝docker-compose(apt常常不能安裝最新版本的docker-compose)
pip install docker-compose
# 4.解決linux下docker的權限問題，將用戶放在docker組里
GROUPNAME=docker
getent group $GROUPNAME 2>&1 >/dev/null || groupadd $GROUPNAME
sudo usermod -aG docker $(whoami)
# 5.重啟
sudo systemctl daemon-reload
sudo systemctl restart docker

使用Docker鏡像

Docker鏡像構建好之后，可以直接運行docker命令啟動鏡像，但是這樣不是最方便的。使用docker-compose搭配Makefile，具體操作如下：首先寫好docker-compose.yml啟動文件，可參考本項目對應的docker-compose.yml，接著，在Makefile里寫常見docker相關命令，我們將應用分為啟動（up）、關閉（down）、進入容器環(huán)境（in）三個需求，Makefile如下：

up:
 docker-compose up -d

down:
 docker-compose down

in:
 docker-compose exec spherical-env bash

本項目鏡像已上傳dockerhub，可以直接使用下列命令下載

docker pull qiangzibro/spherical_image_segmentation
# 或者使用下面命令自己編譯
make build

接著，一鍵完成編譯、啟動

make up #等價于 docker-compose up -d

再通過下列命令便可以進入終端

make in

使用docker-compose logs可以看到notebook對應的網(wǎng)址

2. 數(shù)據(jù)獲取

使用2D-3D-S 數(shù)據(jù)集進行本實驗，該數(shù)據(jù)集提供了來自 2D、2.5D 和 3D 域的各種相互注冊的數(shù)據(jù)，以及實例級語義和幾何注釋。它收集在來自 3 座不同建筑的 6 個大型室內(nèi)區(qū)域。它包含超過 70,000 張 RGB 圖像，以及相應的深度、表面法線、語義注釋、全局 XYZ 圖像（均以常規(guī)和 360° 等距柱狀圖圖像的形式）以及相機信息。它還包括注冊的原始和語義注釋 3D 網(wǎng)格和點云。

數(shù)據(jù)集從開源數(shù)據(jù)集網(wǎng)站格物鈦獲取，這個網(wǎng)站匯總了AI開發(fā)者常見的公開數(shù)據(jù)集，用戶可以對數(shù)據(jù)集進行可視化預覽、在線使用和下載等操作。這里我們不用下載，可以直接通過SDK讀取數(shù)據(jù)集，操作步驟如下：

a. 打開本文對應數(shù)據(jù)集鏈接 https://gas.graviti.cn/dataset/qiangzibro/spherical_segmentation

b. 右上角注冊登錄

c. fork數(shù)據(jù)集

d. 點擊網(wǎng)頁上方開發(fā)者工具，獲取使用SDK所需的AccessKey，獲取到 AccessKey 后，將其存在項目根目錄的gas_key.py里。

KEY = "<Your-Key>"

然后即可以通過AccessKey可以上傳數(shù)據(jù)、讀取數(shù)據(jù)、使用數(shù)據(jù)，靈活對接模型開發(fā)和訓練，與數(shù)據(jù)pipeline快速集成。

e. AccessKey寫入后就可以寫代碼讀取數(shù)據(jù)了。

from PIL import Image

dataset = Dataset("DatasetName", gas)
segment = dataset[0]

for data in segment:
   with data.open() as fp:
       image = Image.open(fp)
       width, height = image.size
       image.show()

3. 方法

使用多種二維CNN方法和球面卷積方法UGSCNN。

其中，二維CNN有三種：UNet、ResNet和FCN；UGSCNN[3]參考自論文《Spherical CNNs on Unstructured Grids》，下面著重看一下UGSCNN的方法。

MeshConv對卷積算子進行定義：

class MeshConv(_MeshConv):
    def __init__(self, in_channels, out_channels, mesh_file, stride=1, bias=True):
        super(MeshConv, self).__init__(in_channels, out_channels, mesh_file, stride, bias)
        pkl = self.pkl
        if stride == 2:
            self.nv_prev = pkl['nv_prev']
            L = sparse2tensor(pkl['L'].tocsr()[:self.nv_prev].tocoo()) # laplacian matrix V->V
            F2V = sparse2tensor(pkl['F2V'].tocsr()[:self.nv_prev].tocoo())  # F->V, #V x #F
        else: # stride == 1
            self.nv_prev = pkl['V'].shape[0]
            L = sparse2tensor(pkl['L'].tocoo())
            F2V = sparse2tensor(pkl['F2V'].tocoo())
        self.register_buffer("L", L)
        self.register_buffer("F2V", F2V)
        
    def forward(self, input):
        # compute gradient
        grad_face = spmatmul(input, self.G)
        grad_face = grad_face.view(*(input.size()[:2]), 3, -1).permute(0, 1, 3, 2) # gradient, 3 component per face
        laplacian = spmatmul(input, self.L)
        identity = input[..., :self.nv_prev]
        grad_face_ew = torch.sum(torch.mul(grad_face, self.EW), keepdim=False, dim=-1)
        grad_face_ns = torch.sum(torch.mul(grad_face, self.NS), keepdim=False, dim=-1)
        grad_vert_ew = spmatmul(grad_face_ew, self.F2V)
        grad_vert_ns = spmatmul(grad_face_ns, self.F2V)

        feat = [identity, laplacian, grad_vert_ew, grad_vert_ns]

        out = torch.stack(feat, dim=-1)
        out = torch.sum(torch.sum(torch.mul(out.unsqueeze(1), self.coeffs.unsqueeze(2)), dim=2), dim=-1)
        out += self.bias.unsqueeze(-1)
        return out

分割網(wǎng)絡基于MeshConv算子構建了一個Unet網(wǎng)絡：

4. 訓練

環(huán)境構建好后只需簡單的幾個命令便可以運行起來

再使用make in成功進入到容器終端

• 基于CNN對網(wǎng)格進行分割

cd cnns
# 基于
./run.sh UNet
# 基于FCN
./run.sh FCN8s
# 基于ResNetDUCHDC
./run.sh ResNetDUCHDC

腳本run.sh解釋

# Model choice
# ResNetDUCHDC,FCN8s,UNet
# Run example
# 1) ./run.sh
# 2) ./run.sh FCN8s
# 3) ./run.sh ResNetDUCHDC
model="${1:-UNet}"
MESHFILES=../data/mesh_files
DATADIR=../data/2d3ds_pano_small/
# create log directory
mkdir -p logs

python train.py \
--batch-size 16 \ # 訓練批量大小
--test-batch-size 16 \ #測試批量大小
--epochs 200 \ # 訓練epoch數(shù)量
--data_folder $DATADIR \
--mesh_folder $MESHFILES \ # 正二十面體網(wǎng)格文件位置
--fold 3 \ # K-fold交叉驗證，k=3。將原始數(shù)據(jù)分成K組(K-Fold)，將每個子集數(shù)據(jù)分別做一次驗證集，其余的K-1組子集數(shù)據(jù)作為訓練集，這樣會得到K個模型。這K個模型分別在驗證集中評估結果，最后的誤差MSE(Mean Squared Error)加和平均就得到交叉驗證誤差。交叉驗證有效利用了有限的數(shù)據(jù)，并且評估結果能夠盡可能接近模型在測試集上的表現(xiàn)，可以做為模型優(yōu)化的指標使用。
--log_dir logs/log_${model}_f16_cv3_rgbd \ # 日志目錄
--decay \ # 學習率衰減
--train_stats_freq 5 \
--model ${model} \ #模型選擇
--in_ch rgbd \ # 輸入數(shù)據(jù)通道
--lr 1e-3 \ # 學習路
--feat 16 #特征層的數(shù)量

cd ugscnn
./run.sh

訓練200個epoch后，可得如下結果：

5. 測試

使用提供的測試腳本test.sh即可進行測試

# 基于UNet
./test.sh UNet
# 基于FCN
./test.sh FCN8s
# 基于ResNetDUCHDC
./test.sh ResNetDUCHDC

測試結果保存在當前目錄下，命名格式為模型名+.npz，將其打開進行結果分析，如下所示。

全景圖實例：

結果：

總結

本文介紹了docker作為環(huán)境構建的知識，介紹幾種基于傳統(tǒng)CNN方法和一種基于球面CNN的方法，并將上述方法在全景數(shù)據(jù)集上完成了分割任務。

數(shù)據(jù)集地址（代碼上傳在數(shù)據(jù)集討論區(qū)）：
https://gas.graviti.cn/dataset/qiangzibro/spherical_segmentation

參考資料