前言

之前接觸到object detection相關的研究，在過程中遇到Dataset不足的問題，因此教授想嘗試看看使用VR環境來產生Dataset來補足資料，並驗證看看對於最後真實資料的detection是否有幫助

而我就負責製作Dataset的生成工具，此篇紀錄我是如何生成VR Dataset

環境配置

Unreal Engine 5
Python 3.9
numpy 1.21.6
AirSim API
wheel 0.42.0
opencv-python

上述的環境配置、安裝可以參考以下文章：

過程

規劃

由於實驗室使用yolo做訓練，因此生成的資料格式會以yolo格式為主

AirSim其實本身有提供生成2D bounding box的API，背後的演算法是由3D bounding box轉換而來，然而AirSim的3D bounding box其實是有歪斜的問題，導致最後得到的2D bounding box可能會有過大的問題

就上述的原因，必須換另一種作法，這邊我想利用segmentation mask作為輔助，透過mask來找到物體的邊界，再依照邊界畫出bounding box

設定生成Classes

首先要標記生成的Classes，可以使用下列的code來獲取classes對應的actor

其中 {object_name}[]* 的意思是，在level中所有以 object_name 為開頭的Actor、Mesh，都會被標記為該class，因此在為level中的物件取名字的時候，務必做好規劃

import airsim
import cv2

client = airsim.VehicleClient()

objects = client.simListSceneObjects(f'{object_name}[\w]*')

接下來要為每個物件做mask的顏色設定，其中255為AirSim的預設mask顏色的數量，因此要做程式保護避免error


for mesh_name in objects:
    if mask_color_cnt + 1 < 255:
        client.simSetSegmentationObjectID(mesh_name, mask_color_cnt + 1, True)

取得Scene、Segmentation mask

接著必須先取得兩張圖片，分別是Scene以及Segmentation mask的圖片，而AirSim得到的資料一開始是string，必須轉換成uint8，且原始圖片要多儲存一份，之後用來畫上bounding box，程式如下：

import airsim
import cv2

# get images response
ori_response, seg_response,  = client.simGetImages(
    [airsim.ImageRequest(camera_name, airsim.ImageType.Scene), airsim.ImageRequest(camera_name, airsim.ImageType.Segmentation)]
)

# get original image from airsim
ori_raw_image = ori_response.image_data_uint8
# get segmentation image from airsim
seg_raw_image = seg_response.image_data_uint8

# trans to uint_8 array form
bbox_img_ary = cv2.imdecode(airsim.string_to_uint8_array(ori_raw_image), cv2.IMREAD_COLOR)
ori_img_ary = cv2.imdecode(airsim.string_to_uint8_array(ori_raw_image), cv2.IMREAD_COLOR)
seg_img_ary = cv2.imdecode(airsim.string_to_uint8_array(seg_raw_image), cv2.IMREAD_COLOR)

ori_img_ary、seg_img_ary輸出後儲存的圖片分別如下：

繪製bounding box

為每個物件標上各自的mask之後，就可以開始生成bounding box，而這邊最簡單的方法就是找到對應顏色的最邊邊的像素點，依照這些像素點的座標去算出四邊形的四個頂點座標

例如：下圖中紅色閃電最邊界的點的座標，就能知道可以剛好框住這個圖案的四邊形頂點分別為：

$ (x_2, y_1), (x_3, y_1) ,(x_2, y_3), (x_3, y_3) $

可以利用下面的API來得到bounding box的資料，x和y代表bounding box左上角的頂點座標，w和h分別代表四邊形的寬和長

import cv2

#contour:mask infomation
x, y, w, h = cv2.boundingRect(contour)

之後將座標轉換為Yolo格式 x, y, w, h

x, y 代表該bndBox的中心座標與圖片寬高的比值，是bndBox歸一化後的中心座標
w, h代表該bndBox的寬高與輸入圖像寬高的比值，是bndBox歸一化後的寬高座標

# Calculate the relative coordinates of the center of the bounding box (normalized coordinates)
center_x = (x + w / 2) / image.shape[1]
center_y = (y + h / 2) / image.shape[0]
relative_width = w / image.shape[1]
relative_height = h / image.shape[0]

最後將所有的bounding box畫到先前預存的原圖上就可以了

結果

最後生成出來的資料如圖：

Github

如果想試用看看生成工具可以到我的github下載，裡面也有我做的VR環境執行檔下載連結

如果有發現bug可以在issue通知我

Kura's BLOG

使用AirSim Plugin生成VR detection dataset

前言

環境配置

過程

規劃

設定生成Classes

取得Scene、Segmentation mask

繪製bounding box

結果

Github