if statement to process frames with float32 dtype

.gitignore

@@ -108,6 +108,15 @@ celerybeat-schedule
 .venv
 venv/
 ENV/
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+myenv
+
 
 # Spyder project settings
 .spyderproject

example_stabilize_frame_16bit.py

@@ -0,0 +1,43 @@
+import os
+import cv2
+from vidstab import VidStab
+import numpy as np
+from tqdm import tqdm
+
+# Init stabilizer and video reader
+stabilizer = VidStab()
+frames = np.load(os.path.join(os.getcwd(), 'vids', 'SEQ_10652.npy'))
+frames = frames.astype(np.float32) / 65535.0
+
+frameidx = 0
+frame_len = frames.shape[2]
+pbar = tqdm(total=frame_len)  # Initialize tqdm with the total number of frames
+
+while frameidx < frame_len:
+    frame = frames[:, :, frameidx]
+
+    # Pass frame to stabilizer even if frame is None
+    stabilized_frame = stabilizer.stabilize_frame(input_frame=frame,
+                                                  border_size=-30,
+                                                  smoothing_window=100)
+
+    # If stabilized_frame is None then there are no frames left to process
+    if stabilized_frame is None:
+        break
+
+    # Display stabilized output
+    cv2.imshow('Stabilized Frame', stabilized_frame)
+    cv2.imshow('Raw Frame', frame)
+
+    key = cv2.waitKey(5)
+    if key == 27:  # Esc key
+        break
+
+    frameidx += 1
+    pbar.update(1)  # Update the progress bar
+
+    if frameidx == frame_len:
+        frameidx = 0
+
+pbar.close()  # Close the progress bar
+cv2.destroyAllWindows()

vidstab/VidStab.py

@@ -162,8 +162,13 @@ def _gen_next_raw_transform(self):
         current_frame_gray = self._resize_frame(current_frame_gray)
 
         # calc flow of movement
-        optical_flow = cv2.calcOpticalFlowPyrLK(self.prev_gray,
-                                                current_frame_gray,
+        prev_gray_uint8 = self.prev_gray
+        current_frame_gray_uint8 = current_frame_gray
+        if current_frame_gray.dtype == np.float32:
+            prev_gray_uint8 = (self.prev_gray * 255).astype(np.uint8)
+            current_frame_gray_uint8 = (current_frame_gray * 255).astype(np.uint8)
+        optical_flow = cv2.calcOpticalFlowPyrLK(prev_gray_uint8,
+                                                current_frame_gray_uint8,
                                                 self.prev_kps, None)
 
         matched_keypoints = vidstab_utils.match_keypoints(optical_flow, self.prev_kps)