From a8afcc81ff9a7d1652702bfb13a1beaef8fae6d7 Mon Sep 17 00:00:00 2001
From: Stephanie Ribet <stephanie.ribet@gmail.com>
Date: Wed, 11 Sep 2024 10:48:10 -0700
Subject: [PATCH] update arina reader

---
 py4DSTEM/io/filereaders/read_arina.py | 43 +++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/py4DSTEM/io/filereaders/read_arina.py b/py4DSTEM/io/filereaders/read_arina.py
index c75679854..ecd9dd5ab 100644
--- a/py4DSTEM/io/filereaders/read_arina.py
+++ b/py4DSTEM/io/filereaders/read_arina.py
@@ -2,7 +2,7 @@
 import hdf5plugin
 import numpy as np
 from py4DSTEM.datacube import DataCube
-from py4DSTEM.utils import bin2D
+from py4DSTEM.preprocess.utils import bin2D
 
 
 def read_arina(
@@ -12,6 +12,8 @@ def read_arina(
     binfactor: int = 1,
     dtype_bin: float = None,
     flatfield: np.ndarray = None,
+    median_filter_masked_pixels_array: np.ndarray = None,
+    median_filter_masked_pixels_kernel: int = 4,
 ):
     """
     File reader for arina 4D-STEM datasets
@@ -79,6 +81,8 @@ def read_arina(
             array_3D,
             binfactor,
             correction_factors,
+            median_filter_masked_pixels_array,
+            median_filter_masked_pixels_kernel,
         )
 
     if f.__bool__():
@@ -95,23 +99,50 @@ def read_arina(
         )
     )
 
+    if median_filter_masked_pixels_array is not None and binfactor == 1:
+        datacube = datacube.median_filter_masked_pixels(
+            median_filter_masked_pixels_array, median_filter_masked_pixels_kernel
+        )
+
     return datacube
 
 
-def _processDataSet(dset, start_index, array_3D, binfactor, correction_factors):
+def _processDataSet(
+    dset,
+    start_index,
+    array_3D,
+    binfactor,
+    correction_factors,
+    median_filter_masked_pixels_array,
+    median_filter_masked_pixels_kernel,
+):
     image_index = start_index
     nimages_dset = dset.shape[0]
 
+    if median_filter_masked_pixels_array is not None and binfactor != 1:
+        from py4DSTEM.preprocess import median_filter_masked_pixels_2D
+
     for i in range(nimages_dset):
         if binfactor == 1:
             array_3D[image_index] = np.multiply(
                 dset[i].astype(array_3D.dtype), correction_factors
             )
+
         else:
-            array_3D[image_index] = bin2D(
-                np.multiply(dset[i].astype(array_3D.dtype), correction_factors),
-                binfactor,
-            )
+            if median_filter_masked_pixels_array is not None:
+                array_3D[image_index] = bin2D(
+                    median_filter_masked_pixels_2D(
+                        np.multiply(dset[i].astype(array_3D.dtype), correction_factors),
+                        median_filter_masked_pixels_array,
+                        median_filter_masked_pixels_kernel,
+                    ),
+                    binfactor,
+                )
+            else:
+                array_3D[image_index] = bin2D(
+                    np.multiply(dset[i].astype(array_3D.dtype), correction_factors),
+                    binfactor,
+                )
 
         image_index = image_index + 1
     return image_index