omni3d_utils.py

omni3d_utils.py

def correct_obj_bbox(obj, h, w, expand_rate=0.05):
    if obj['bbox2D_tight'][0] >= 0.0:
        obj['bbox2D_tight'] = expand_bbox(obj['bbox2D_tight'], h, w, expand_rate)
    elif obj['bbox2D_trunc'][0] >= 0.0:
        obj['bbox2D_tight'] = expand_bbox(obj['bbox2D_trunc'], h, w, expand_rate)
    else:
        obj['bbox2D_tight'] = expand_bbox(obj['bbox2D_proj'], h, w, -expand_rate)
    return obj


def create_data_mapping(annot):
    annotations = annot['annotations']
    mapping = {}
    for a in annotations:
        if a['image_id'] not in mapping: mapping[a['image_id']] = []
        mapping[a['image_id']].append(a)
    return mapping


if __name__ == '__main__':
    for data_name in DATASETS:
        annot = None
        for s in args.splits:
            annot_file = os.path.join(args.data_path, 'Omni3D', f'{data_name}_{s}.json')
            _annot = json.load(open(annot_file))
            if annot is None:
                annot = _annot
            else:
                annot['images'] = annot['images'] + _annot['images']
                annot['annotations'] = annot['annotations'] + _annot['annotations']
        data_mapping = create_data_mapping(annot)