lint

loicmagne · Jun 17, 2024 · d28a9df · d28a9df
1 parent 9119dd9
commit d28a9df
Showing 1 changed file with 18 additions and 10 deletions.
diff --git a/align-rec/aligner.py b/align-rec/aligner.py
@@ -4,6 +4,7 @@
 import polars as pl
 import numpy as np
 
+
 def align(
     path_video: str,
     path_inputs: str,
@@ -27,23 +28,29 @@ def align(
             ret, frame = cap.read()
             if not ret:
                 break
-
-            t_micro = cap.get(cv2.CAP_PROP_POS_MSEC) * 1000 # Frame timestamp in microseconds
-            while inputs['time'][i] < t_micro: # Find the closest timestamp in inputs
+
+            t_micro = (
+                cap.get(cv2.CAP_PROP_POS_MSEC) * 1000
+            )  # Frame timestamp in microseconds
+            while inputs["time"][i] < t_micro:  # Find the closest timestamp in inputs
                 i += 1
 
             # Pick the closest timestamp between the two neighbors
-            prev = inputs['time'][i-1]
-            next = inputs['time'][i]
-            row = inputs.row(i-1, named=True) if abs(prev - t_micro) < abs(next - t_micro) else inputs.row(i, named=True)
+            prev = inputs["time"][i - 1]
+            next = inputs["time"][i]
+            row = (
+                inputs.row(i - 1, named=True)
+                if abs(prev - t_micro) < abs(next - t_micro)
+                else inputs.row(i, named=True)
+            )
             inputs_aligned.append(row)
-            
-            ERRORS.append(abs(row['time'] - t_micro))
+
+            ERRORS.append(abs(row["time"] - t_micro))
             COUNT_MANUAL += 1
             pbar.update(1)
     cap.release()
     inputs_aligned = pl.DataFrame(inputs_aligned)
-    inputs_aligned.write_parquet(Path(path_inputs).with_suffix('.parquet'))
+    inputs_aligned.write_parquet(Path(path_inputs).with_suffix(".parquet"))
 
     print(f"Number of frames (OpenCV): {COUNT_OPENCV}")
     print(f"Number of frames (Manual): {COUNT_MANUAL}")
@@ -52,6 +59,7 @@ def align(
     print(f"Min error: {np.min(ERRORS)}")
     print(inputs_aligned)
 
+
 def main_cli():
     import argparse
 
@@ -60,4 +68,4 @@ def main_cli():
     parser.add_argument("inputs", help="Path to inputs .csv file")
     args = parser.parse_args()
 
-    align(args.video, args.inputs)
+    align(args.video, args.inputs)