Files
hivecore_robot_vision/tools/data_tools.py
liangyuxuan ea1a985b44 add tools
2025-11-26 10:30:55 +08:00

55 lines
2.4 KiB
Python

import argparse
import json
import transforms3d as tfs
def prepare_data(_data:list[float], save_path: str = "camera_pose_data.json", format: str = "quat_xyzw", mm2m: bool = False):
new_list = []
if format == "quat_wxyz":
for i in range(0, len(_data), 7):
if not mm2m:
new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 3], _data[i + 4], _data[i + 5], _data[i + 6]]
else:
new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5], _data[i + 6]]
new_list.extend(new)
elif format == "quat_xyzw":
for i in range(0, len(_data), 7):
if not mm2m:
new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
else:
new = [_data[i] / 1e3, _data[i + 1] / 1e3, _data[i + 2] / 1e3, _data[i + 6], _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
new_list.extend(new)
elif format == "euler":
for i in range(0, len(_data), 6):
rw, rx, ry, rz = tfs.euler.euler2quat(_data[i + 3], _data[i + 4], _data[i + 5], axes='sxyz')
if not mm2m:
new = [_data[i], _data[i + 1], _data[i + 2], rw, rx, ry, rz]
else:
new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, rw, rx, ry, rz]
new_list.extend(new)
elif format == "rotvertor":
for i in range(0, len(_data), 6):
# tfs.quaternions.rotvec2quat
if not mm2m:
new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 3], _data[i + 4], _data[i + 5]]
else:
new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
new_list.extend(new)
else: pass
with open(save_path, 'w') as f: json.dump(new_list, f, indent=4)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--data_path", type=str, required=True)
parser.add_argument("--save_path", type=str, required=True)
parser.add_argument("--format", type=str, default="quat_xyzw")
parser.add_argument("--mm2m", type=bool, default=False)
args = parser.parse_args()
with open(args.data_path, "r") as f:
_data = json.load(f)
prepare_data(_data, args.save_path, args.format, args.mm2m)