AirSim图像数据集收集

AirSim图像数据集收集

---

前言

从该文开始为视觉目标检测做准备，本文介绍一种半自动收集数据集的方法。

---

一、本文实现的效果

以下的图像数据集均为程序自动收集的，我们通过手动控制无人机切换不同位置以及视角来收集具有多样性的图像数据集。

二、环境依赖

1.基本说明

如果跟着之前的文章做实验，那么恭喜你，这次环境依赖并没有增加。如果你没有看过之前的文章，那么在进行实验之前需要安装pygame库以及opencv-python库。其中pygame库用来做基本的可视化以及实现多按键控制，opencv库用来对AirSim的图像API返回的图像数据进行解码，以及保存图像，写入特定文件夹。

2.安装依赖

安装指令如下：

pip3 install pygame
pip3 install opencv-python

---

三、程序实现

1.基本说明

当前能够较实时读取的图像分辨率都比较低，事实上，在AirSim中，读取1920x1080分辨率的图像的时候，每帧时间高达300+ms，对视觉集群来说是基本不可行的，而且高分辨率图像给检测带来了较大挑战，我们也没有必要使用到如此高的分辨率。在多次实验测试后，我们决定使用480P的图像进行视觉集群，在单目距离估计中，对于DJI M200模型的无人机对象，如下图所示，该分辨率下有效检测距离达到10米以上，是完全可以满足我们的需求的。

基本的程序流程如下所示:

要注意的是，图像的保存名称应该是动态变化的，否则新图像会不断覆盖旧图像。为了解决这个问题，我们应该对每次的新图像起新名称。容易想到的两种解决方案是使用时间戳或者使用文件夹中文件总数，第二种有覆盖的风险，这里，我们采用时间戳的方法，基本使用如下所示。

import datetime

save_name = images_save_path + str(datetime.datetime.now().strftime('date%w_%H_%M_%S')) + ".jpg"

2.implementation

代码在进阶版的键盘控制基础上增加自动保存数据集的功能，默认的控制可能并不符合自己的手感，可以随意调节，另外在程序中需要修改images_save_path变量。

import sys
import cv2
import time
import airsim
import pygame
import datetime

# >------>>>  pygame settings   <<<------< #
pygame.init()
screen = pygame.display.set_mode((480, 360))
pygame.display.set_caption('keyboard ctrl @FPV')
screen.fill((0, 0, 0))
# >------>>>   --------------------------------    <<<------< #

# >------>>>  AirSim settings   <<<------< #
# 这里改为你要控制的无人机名称(settings文件里面设置的)
base_name = "Drone"
vehicle_index = [0, 1]
now_index = vehicle_index[0]

AirSim_client = airsim.MultirotorClient()
AirSim_client.confirmConnection()

AirSim_client.enableApiControl(True, vehicle_name=base_name + str(vehicle_index[0]))
AirSim_client.enableApiControl(True, vehicle_name=base_name + str(vehicle_index[1]))

AirSim_client.armDisarm(True, vehicle_name=base_name + str(vehicle_index[0]))
AirSim_client.armDisarm(True, vehicle_name=base_name + str(vehicle_index[1]))

AirSim_client.takeoffAsync(vehicle_name=base_name + str(vehicle_index[0]))
AirSim_client.takeoffAsync(vehicle_name=base_name + str(vehicle_index[1])).join()

image_types = {
            "scene": airsim.ImageType.Scene,
            "depth": airsim.ImageType.DepthVis,
            "seg": airsim.ImageType.Segmentation,
            "normals": airsim.ImageType.SurfaceNormals,
            "segmentation": airsim.ImageType.Segmentation,
            "disparity": airsim.ImageType.DisparityNormalized
        }
# >------>>>   --------------------------------    <<<------< #

# >------>>>  manual controller settings   <<<------< #
# 基础的控制速率
base_rate = 0.2
# 悬停时的油门
base_throttle = 0.55
# 设置临时加速比例
speedup_ratio = 4.0
# 用来设置临时加速
speedup_flag = False

change_time = 0.0
# 防止来回切换控制对象
enable_change = True
# 用来标记当前是否进行键盘控制
control_iteration = False
# >------>>>   --------------------------------    <<<------< #


# >------>>>  datasets collect settings   <<<------< #
# 1.5s 保存一帧
start_time = 0.0
save_period = 1.5
save_all_number = 0
# 需要把保存路径改成自己的文件夹
images_save_path = "../data/datasets/temp/"


while True:

    pitch_rate = 0.0
    yaw_rate = 0.0
    roll_rate = 0.0
    throttle = base_throttle
    control_iteration = False

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            sys.exit()

    scan_wrapper = pygame.key.get_pressed()

    # 按下空格键加速10倍
    if scan_wrapper[pygame.K_SPACE]:
        scale_ratio = speedup_ratio
    else:
        scale_ratio = speedup_ratio / speedup_ratio

    # 切换两秒后方可再次切换
    if time.time() - change_time > 2:
        enable_change = True

    # 需要加上enable_change变量判断，否则会来回切换
    if scan_wrapper[pygame.K_LCTRL] and scan_wrapper[pygame.K_c] and enable_change:
        enable_change = False
        change_time = time.time()
        if now_index == vehicle_index[0]:
            now_index = vehicle_index[1]
        else:
            now_index = vehicle_index[0]
        print(f"change to drone{now_index} ···")
        time.sleep(0.2)

    # 根据 'A' 和 'D' 按键来设置偏航速率变量
    if scan_wrapper[pygame.K_a] or scan_wrapper[pygame.K_d]:
        control_iteration = True
        yaw_rate = (scan_wrapper[pygame.K_a] - scan_wrapper[pygame.K_d]) * scale_ratio * base_rate

    # 根据 'UP' 和 'DOWN' 按键来设置pitch轴速率变量(NED坐标系，pitch为机头向前)
    if scan_wrapper[pygame.K_UP] or scan_wrapper[pygame.K_DOWN]:
        control_iteration = True
        pitch_rate = (scan_wrapper[pygame.K_UP] - scan_wrapper[pygame.K_DOWN]) * scale_ratio * base_rate

    # 根据 'LEFT' 和 'RIGHT' 按键来设置roll轴速率变量(NED坐标系，roll为正右方)
    if scan_wrapper[pygame.K_LEFT] or scan_wrapper[pygame.K_RIGHT]:
        control_iteration = True
        roll_rate = -(scan_wrapper[pygame.K_LEFT] - scan_wrapper[pygame.K_RIGHT]) * scale_ratio * base_rate

    # 根据 'W' 和 'S' 按键来设置z轴速率变量(NED坐标系，z轴向上为负，油门为正数)
    if scan_wrapper[pygame.K_w] or scan_wrapper[pygame.K_s]:
        control_iteration = True
        throttle = base_throttle + (scan_wrapper[pygame.K_w] - scan_wrapper[pygame.K_s]) * scale_ratio * base_rate

    # 速率需要限幅
    if pitch_rate > 1.0:
        pitch_rate = 1.0
    elif pitch_rate < -1.0:
        pitch_rate = -1.0

    if yaw_rate > 1.0:
        yaw_rate = 1.0
    elif yaw_rate < -1.0:
        yaw_rate = -1.0

    if roll_rate > 1.0:
        roll_rate = 1.0
    elif roll_rate < -1.0:
        roll_rate = -1.0

    if throttle > 1.0:
        throttle = 1.0
    elif throttle < 0.0:
        throttle = 0.0

    # 获取正前方的一帧图像
    temp_image = AirSim_client.simGetImage('0', image_types["scene"], vehicle_name=base_name + str(now_index))

    if temp_image is None:
        print("Warning: Failed to read a frame!! ")
        pygame.quit()

    else:
        pass

    # 将图像进行解码，变成像素值为0-255的范围，保存路径自行修改
    image = cv2.imdecode(airsim.string_to_uint8_array(temp_image), cv2.IMREAD_COLOR)
    cv2.imwrite("visual.png", image)

    # 不控制的时候保持悬停，也可以用于刹车
    if control_iteration:
        # 设置速率控制以及设置偏航控制
        AirSim_client.moveByRollPitchYawrateThrottleAsync(pitch=pitch_rate, roll=roll_rate, yaw_rate=yaw_rate,
                                                          throttle=throttle, duration=0.05,
                                                          vehicle_name=base_name + str(now_index))
    else:
        # 保持当前位置
        AirSim_client.hoverAsync(vehicle_name=base_name + str(now_index))

    # 自动生成的图片名称
    save_name = images_save_path + str(datetime.datetime.now().strftime('date%w_%H_%M_%S')) + ".jpg"
    if time.time() - start_time >= save_period:
        success = cv2.imwrite(save_name, image)
        if success:
            save_all_number += 1
            print("\rsaving the {:4d}  frames...".format(save_all_number), end='')
        else:
            print(f"Warning: Failed to save {save_name}")
        start_time = time.time()

    # 利用pygame库加载保持的第一视角图像，
    screen_image = pygame.image.load("visual.png")
    # 图像坐标系，左上角为(0, 0)，在此放置图片
    screen.blit(screen_image, (0, 0))
    pygame.display.flip()
    pygame.display.update()

    # press 'Esc' to quit
    if scan_wrapper[pygame.K_ESCAPE]:
        pygame.quit()
        sys.exit()

总结

该文提供了一种半自动收集数据集的方法，在控制无人机进行飞行时程序自动定时收集数据集，事实上，在实际应用时应该采用快捷键启动数据集收集与停止数据集收集，不过在前两篇博文的基础上，应该自己就能轻松想到解决方案，这个工作留给读者自己去做。

下篇预告

随着代码功能块分散，程序可读性会越来越差，下篇博文将开始进行功能封装，并且开始搭建yolov5的深度学习目标检测算法。