diff --git a/bbox_2d.py b/bbox_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..f039ae370c38ba61b32b6da0f5bf534bd48dcd2b --- /dev/null +++ b/bbox_2d.py @@ -0,0 +1,49 @@ +import math +import numpy as np +import json + + +def bbox_2d(data_sample): + + center = data_sample['bbox']['center'] + size = data_sample['bbox']['size'] + bbox_res = [] + top_left_x, top_left_y = center[0] - (size[0] / 2), center[1] - (size[1] / 2) + bbox_res.append(top_left_x) + bbox_res.append(top_left_y) + + bottom_right_x, bottom_right_y = center[0] + (size[0] / 2), center[1] + (size[1] / 2) + bbox_res.append(bottom_right_x) + bbox_res.append(bottom_right_y) + + return bbox_res + + +def generate_2d_bbox(data_name): + + cont1, cont2, cont3 = 0, 0, 0 + num_arr = 0 + for p in range(4995): + with open(f"{data_name}/Bbox_2d/{p}.json", 'r') as f: + data = json.load(f) + + for i in range(len(data)): + if data[i]['id'] == 4: + cont1 += 1 + bbox = bbox_2d(data[i]) + np.savetxt(f'{data_name}/Bbox/Banana/{p}.txt', np.array(bbox).reshape((1, 4))) # save + + elif data[i]['id'] == 5: + cont2 += 1 + bbox = bbox_2d(data[i]) + np.savetxt(f'{data_name}/Bbox/Orange/{p}.txt', np.array(bbox).reshape((1, 4))) # save + + elif data[i]['id'] == 6: + cont3 += 1 + bbox = bbox_2d(data[i]) + np.savetxt(f'{data_name}/Bbox/Pear/{p}.txt', np.array(bbox).reshape((1, 4))) # save + + else: + continue + print(cont1, cont2, cont3) +