Add a script for modifying data.jsonl files

Author: oseiskar

modify-data-jsonl.py

@@ -0,0 +1,72 @@
+"""
+Modify Spectacular AI data JSONL files
+
+Input JSONL in read from stdin an output written to stdout
+"""
+import json
+import sys
+
+def modify_jsonl_sorted(in_f, out_f, func):
+    lines = []
+    for line in in_f:
+        d = func(json.loads(line))
+        if d is not None:
+            lines.append(d)
+    for l in sorted(lines, key=lambda d: d.get('time', -1e100)):
+        out_f.write(json.dumps(l)+'\n')
+
+def modify_jsonl_stream(in_f, out_f, func):
+    for line in in_f:
+        d = func(json.loads(line))
+        if d is not None:
+            out_f.write(json.dumps(d)+'\n')
+    return out_f
+
+def shift_imu_time(data, delta_t):
+    if 'sensor' in data:
+        data['time'] += delta_t
+    return data
+
+def shift_frame_times(data, delta_t):
+    if 'frames' in data:
+        data['time'] += delta_t
+        for cam in data['frames']:
+            if 'time' in cam: cam['time'] += delta_t
+    return data
+
+def drop_zero_imu_samples(data):
+    if 'sensor' in data and all([v == 0.0 for v in data['sensor']['values']]): return None
+    return data
+
+if __name__ == '__main__':
+    import argparse
+    p = argparse.ArgumentParser(__doc__)
+    subs = p.add_subparsers(dest='command')
+
+    sub_shift_imu = subs.add_parser('shift_imu_time', help='shift IMU timestamps by given delta_t')
+    sub_shift_imu.add_argument('delta_t', help='timestamp shift in seconds', type=float)
+
+    sub_shift_frames = subs.add_parser('shift_frame_times', help='shift frame timestamps by given delta_t')
+    sub_shift_frames.add_argument('delta_t', help='timestamp shift in seconds', type=float)
+
+    sub_drop_zero_imus = subs.add_parser('drop_zero_imu_samples', help='remove all-zero IMU samples')
+
+    p.add_argument('-s', '--sort', action='store_true', help='sort the resulting file based on timestamp')
+    
+    args = p.parse_args()
+    f_in = sys.stdin
+    f_out = sys.stdout
+
+    if args.command == 'shift_imu_time':
+        func = lambda x: shift_imu_time(x, args.delta_t)
+    elif args.command == 'shift_frame_times':
+        func = lambda x: shift_frame_times(x, args.delta_t)
+    elif args.command == 'drop_zero_imu_samples':
+        func = drop_zero_imu_samples
+    else:
+        assert(False)
+
+    if args.sort:
+        modify_jsonl_sorted(f_in, f_out, func)
+    else:
+        modify_jsonl_stream(f_in, f_out, func)