HolisticDiffuScene/visualize_bbox.py at develop · fangchuan/HolisticDiffuScene · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import os
import json
import argparse

import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation as R

from misc.utils import get_corners_of_bb3d_no_index, project_3d_points_to_2d, parse_camera_info, matrix_to_euler_angles
from misc.equirect_projection import vis_objs3d
from dataset.metadata import INVALID_SCENES_LST, INVALID_ROOMS_LST, OBJECT_LABEL_IDS

TARGET_ROOM_TYPE_LST = ['living room', 'bedroom', 'dining room', 'kitchen', 'bathroom', 'study']


def visualize_bbox(dataset_folderpath, output_folderpath):

    SCENE_LST = ['scene_%05d' % i for i in range(0, 3500) if ('scene_%05d' % i) not in INVALID_SCENES_LST]

    for scene_id in SCENE_LST:
        scene_img_path = os.path.join(dataset_folderpath, scene_id, "2D_rendering")

        room_type_lst = None
        scene_anno_3d_filepath = os.path.join(dataset_folderpath, scene_id, "annotation_3d.json")
        if not os.path.isfile(scene_anno_3d_filepath):
            INVALID_SCENES_LST.append(scene_id)
            continue
        else:
            scene_anno_3d_dict = json.load(open(scene_anno_3d_filepath, 'r'))
            room_type_lst = scene_anno_3d_dict['semantics']

        for room_id in np.sort(os.listdir(scene_img_path)):
            room_id_str = scene_id + '_' + room_id
            if room_id_str in INVALID_ROOMS_LST:
                continue

            room_path = os.path.join(scene_img_path, room_id, "panorama")
            if not os.path.exists(room_path):
                continue
            room_bbox_3d_path = os.path.join(room_path, 'full', 'bbox_3d.json')
            with open(room_bbox_3d_path, 'r') as file:
                annos = json.load(file)

            id2index = dict()
            for index, object in enumerate(annos):
                id2index[object.get('ID')] = index

            room_type_str = 'undefined'
            if room_type_lst is not None:
                for rt in room_type_lst:
                    if rt['ID'] == int(room_id):
                        room_type_str = rt['type']
                        break
            if room_type_str != 'undefined':
                os.makedirs(os.path.join(output_folderpath, room_type_str), exist_ok=True)

            rgb_img = cv2.imread(os.path.join(room_path, 'full', 'rgb_rawlight.png'))
            rgb_img = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2RGB)
            height, width, _ = rgb_img.shape

            instance_img = cv2.imread(os.path.join(room_path, 'full', 'instance.png'), cv2.IMREAD_UNCHANGED)
            depth_img = cv2.imread(os.path.join(room_path, 'full', 'depth.png'), cv2.IMREAD_UNCHANGED)

            cam_position = np.loadtxt(os.path.join(room_path, 'camera_xyz.txt'))
            cam_position = cam_position
            # print(f'cam_position: {cam_position}')

            obj_bbox_lst = []
            # skip background
            for index in np.unique(instance_img)[:-1]:
                # for each instance in current image
                bbox = annos[id2index[index]]

                if bbox['label'] not in OBJECT_LABEL_IDS.values():
                    continue

                basis = np.array(bbox['basis'])
                coeffs = np.array(bbox['coeffs'])
                centroid = np.array(bbox['centroid'])

                obj_bbox_dict = {}
                obj_bbox_dict['rotations'] = basis.tolist()
                obj_bbox_dict['centroid'] = list(centroid)
                obj_bbox_dict['dimensions'] = list(coeffs)
                obj_bbox_dict['name'] = bbox['label']

                # obj_bbox_dict['angles'] = R.from_matrix(basis).as_euler('xyz', degrees=False).tolist()
                obj_bbox_dict['angles'] = matrix_to_euler_angles(basis).tolist()
                if bbox['label'] == 'door':
                    print(f'basis of door: {basis}')
                    print(f'angles of door: {obj_bbox_dict["angles"]}')
                obj_bbox_dict['center'] = list((centroid - cam_position) * 0.001)
                obj_bbox_dict['size'] = list(coeffs * 0.001 * 2)
                obj_bbox_lst.append(obj_bbox_dict)

            if room_type_str != 'undefined':
                anno_img = vis_objs3d(image=rgb_img,
                                      v_bbox3d=obj_bbox_lst,
                                      camera_position=cam_position,
                                      b_show_axes=False,
                                      b_show_centroid=False,
                                      b_show_bbox3d=True,
                                      b_show_info=True,
                                      thickness=2)
                output_img_filepath = os.path.join(output_folderpath, room_type_str, room_id_str + '_bbox.png')
                print(f'save visualization for object bbox annotation of {room_id_str}')
                cv2.imwrite(output_img_filepath, anno_img)
                save_obj_bbox_filepath = os.path.join(output_folderpath, room_type_str, room_id_str + '_bbox.json')
                obj_bbox_dicts = {}
                obj_bbox_dicts['objects'] = obj_bbox_lst
                json.dump(obj_bbox_dicts, open(save_obj_bbox_filepath, 'w'), indent=4)


def parse_args():
    parser = argparse.ArgumentParser(description="Structured3D 3D Bounding Box Visualization")
    parser.add_argument("--dataset_path",
                        default="/data/dataset/Structured3D/Structured3D/",
                        help="raw dataset path",
                        metavar="DIR")
    parser.add_argument("--debug_path",
                        default="/data/dataset/Structured3D/preprocessed/debug_annotations",
                        help="debug folder path for object bbox annotations",
                        metavar="DIR")
    return parser.parse_args()


def main():
    args = parse_args()
    dataset_folderpath = args.dataset_path
    debug_folderpath = args.debug_path

    if not os.path.exists(dataset_folderpath):
        raise ValueError("Dataset folder does not exist!")
        exit(-1)

    if not os.path.exists(debug_folderpath):
        os.makedirs(debug_folderpath)

    visualize_bbox(dataset_folderpath, debug_folderpath)


if __name__ == "__main__":
    main()