Make module I/Os a bit decoupled

miniscope_io/processing/video.py

@@ -83,18 +83,18 @@ def __del__(self):
         self.release()
 
 def gen_freq_mask(
-        img: np.ndarray,
+        width: int = 200,
+        height: int = 200,
         show_mask: bool = True
         ) -> np.ndarray:
     """
     Generate a mask to filter out horizontal and vertical frequencies.
     A central circular region can be removed to allow low frequencies to pass.
     """
-    rows, cols = img.shape
-    crow, ccol = rows // 2, cols // 2
+    crow, ccol = height // 2, width // 2
 
     # Create an initial mask filled with ones (pass all frequencies)
-    mask = np.ones((rows, cols), np.uint8)
+    mask = np.ones((height, width), np.uint8)
 
     # Define band widths for vertical and horizontal suppression
     vertical_band_width = 2
@@ -108,7 +108,7 @@ def gen_freq_mask(
 
     # Define the radius of the circular region to retain at the center
     radius = 6
-    y, x = np.ogrid[:rows, :cols]
+    y, x = np.ogrid[:height, :width]
     center_mask = (x - ccol) ** 2 + (y - crow) ** 2 <= radius ** 2
 
     # Restore the center circular area to allow low frequencies to pass
@@ -128,15 +128,21 @@ class FrameProcessor:
     """
     A class to process video frames.
     """
-    def __init__(self, img, block_size=400):
+    def __init__(self,
+                 height: int,
+                 width: int,
+                 buffer_size: int=5032,
+                 block_size: int=32
+                 ):
         """
         Initialize the FrameProcessor object.
+        Block size/buffer size will be set by dev config later.
         """
-        self.mask = gen_freq_mask(img)
-        self.previous_frame = img
+        self.height = height
+        self.width = width
+        self.buffer_size = buffer_size
         self.block_size = block_size
 
-
     def is_block_noisy(
             self, block: np.ndarray, prev_block: np.ndarray, noise_threshold: float
             ) -> bool:
@@ -150,32 +156,34 @@ def is_block_noisy(
         # Consider noisy if the mean difference exceeds the threshold
         return mean_diff > noise_threshold
 
-    def process_frame(
-            self, current_frame: np.ndarray, noise_threshold: float
+    def patch_noisy_buffer(
+            self,
+            current_frame: np.ndarray,
+            previous_frame: np.ndarray,
+            noise_threshold: float
             ) -> np.ndarray:
         """
         Process the frame, replacing noisy blocks with those from the previous frame.
-
-        Args:
-            current_frame: The current frame to process.
-            noise_threshold: The threshold for determining noisy blocks.
         """
         processed_frame = np.copy(current_frame)
         h, w = current_frame.shape
 
         for y in range(0, h, self.block_size):
             for x in range(0, w, self.block_size):
                 current_block = current_frame[y:y+self.block_size, x:x+self.block_size]
-                prev_block = self.previous_frame[y:y+self.block_size, x:x+self.block_size]
+                prev_block = previous_frame[y:y+self.block_size, x:x+self.block_size]
 
                 if self.is_block_noisy(current_block, prev_block, noise_threshold):
                     # Replace current block with previous block if noisy
                     processed_frame[y:y+self.block_size, x:x+self.block_size] = prev_block
 
-        self.previous_frame = processed_frame
         return processed_frame
 
-    def remove_stripes(self, img):
+    def remove_stripes(
+            self,
+            img: np.ndarray,
+            mask: np.ndarray
+            )-> np.ndarray:
         """Perform FFT/IFFT to remove horizontal stripes from a single frame."""
         f = np.fft.fft2(img)
         fshift = np.fft.fftshift(f)
@@ -194,7 +202,7 @@ def remove_stripes(self, img):
         logger.debug(f"FFT shape: {fshift.shape}")
 
         # Apply mask and inverse FFT
-        fshift *= self.mask
+        fshift *= mask
         f_ishift = np.fft.ifftshift(fshift)
         img_back = np.fft.ifft2(f_ishift)
         img_back = np.abs(img_back)
@@ -215,17 +223,34 @@ def remove_stripes(self, img):
     frames = []
     index = 0
 
+    processor = FrameProcessor(
+        height=200,
+        width=200,
+    )
+
+    freq_mask = gen_freq_mask(
+        width=200,
+        height=200,
+        show_mask=True
+        )
     try:
         for frame in reader.read_frames():
             gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-            if index == 0:
-                processor = FrameProcessor(gray_frame)
+            index += 1
             if index > 100:
                 break
-            index += 1
             logger.info(f"Processing frame {index}")
-            processed_frame = processor.process_frame(gray_frame, noise_threshold=10)
-            filtered_frame = processor.remove_stripes(processed_frame)
+            if index == 1:
+                previous_frame = gray_frame
+            processed_frame = processor.patch_noisy_buffer(
+                gray_frame,
+                previous_frame,
+                noise_threshold=30
+                )
+            filtered_frame = processor.remove_stripes(
+                img=processed_frame,
+                mask=freq_mask
+                )
             frames.append(filtered_frame)
 
         # Display the results for visualization