diff --git a/.gitignore b/.gitignore index 33dee03..9d59044 100644 --- a/.gitignore +++ b/.gitignore @@ -108,6 +108,15 @@ celerybeat-schedule .venv venv/ ENV/ +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +myenv + # Spyder project settings .spyderproject diff --git a/example_stabilize_frame_16bit.py b/example_stabilize_frame_16bit.py new file mode 100644 index 0000000..cb8c4fa --- /dev/null +++ b/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() diff --git a/vids/SEQ_10652.npy b/vids/SEQ_10652.npy new file mode 100644 index 0000000..0a05e49 Binary files /dev/null and b/vids/SEQ_10652.npy differ diff --git a/vidstab/VidStab.py b/vidstab/VidStab.py index 67a5033..c42e649 100644 --- a/vidstab/VidStab.py +++ b/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)