AI2Thor环境问题

低版本AI2Thor环境取obj的bbx

"""AI2Thor版本1.0.1"""
from ai2thor.controller import Controller


# 初始化AI2-THOR控制器
controller = Controller()

# 启动AI2-THOR仿真环境
controller.start()

# 定义要遍历的房间
rooms = [f'FloorPlan{i}' for i in range(1, 31)] + \
        [f'FloorPlan{i}' for i in range(201, 231)]   # 添加要遍历的房间名称

# 设置每个房间的物体bounding box保存目录
output_dir = 'image_data'
os.makedirs(output_dir, exist_ok=True)

# 循环遍历每个房间
for room in rooms:
    # 重置仿真环境并设置当前房间
    controller.reset(room)

    controller.step(dict(action='Initialize', gridSize=0.25, fieldOfView=90, renderObjectImage=True))        # 在这里说明renderObjectImage=True，后面才能输出2D检测结果
    file_path = f"/dataset/{room}/images.hdf5"
    # 打开HDF5文件
    with h5py.File(file_path, 'r') as file:
        for dataset_name in file:
            positions = dataset_name.split("|")          # 得到分隔后的当前房间的每个坐标点
            event = controller.step(dict(action='TeleportFull', x=positions[0], y=0.9009992, z=positions[1], rotation={'y': positions[2]}, horizon=positions[3]))
            bbx_dict = event.class_detections2D                     # 获得检测数据
                for i in list(event.class_detections2D.keys()):
                    # 将内部的NumPy数组转换为Python列表
                    bbx = [item.tolist() for item in event.class_detections2D[i]]
                    bbx_dict[i] = bbx

     # 保存就行了