Added support for fits.gz files

ryanhausen · Jul 12, 2024 · 2cb2ccc · 2cb2ccc
1 parent 329ccde
commit 2cb2ccc
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/fitsmap/convert.py b/fitsmap/convert.py
@@ -64,7 +64,7 @@
 
 Shape = Tuple[int, int]
 
-IMG_FORMATS = ["fits", "jpg", "png"]
+IMG_FORMATS = [".fits", ".fits.gz", ".jpg", ".png"]
 CAT_FORMAT = ["cat"]
 IMG_ENGINE_PIL = "PIL"
 IMG_ENGINE_MPL = "MPL"
@@ -156,9 +156,7 @@ def get_array(file_location: str) -> pa.PaddedArray:
         A numpy array representing the image.
     """
 
-    _, ext = os.path.splitext(file_location)
-
-    if ext == ".fits":
+    if file_location.endswith(".fits") or file_location.endswith(".fits.gz"):
         array = fits.getdata(file_location).astype(np.float32)
         shape = array.shape
         if len(shape) != 2:
@@ -190,7 +188,7 @@ def filter_on_extension(
 
     return list(
         filter(
-            lambda s: (os.path.splitext(s)[1][1:] in extensions)
+            lambda s: any((s.endswith(e) for e in extensions))
             and not neg_predicate(s),
             files,
         )
@@ -455,7 +453,7 @@ def tile_img(
 
     # if we're using matplotlib we need to instantiate the matplotlib objects
     # before we pass them to ray
-    image_engine = IMG_ENGINE_MPL if file_location.endswith(".fits") else IMG_ENGINE_PIL
+    image_engine = IMG_ENGINE_MPL if (file_location.endswith(".fits") or file_location.endswith(".fits.gz")) else IMG_ENGINE_PIL
     if image_engine == IMG_ENGINE_MPL:
         image_norm = norm_kwargs.get(os.path.basename(file_location), norm_kwargs)
         mpl_norm, mpl_cmap = build_mpl_objects(array.array, image_norm)

diff --git a/fitsmap/utils.py b/fitsmap/utils.py
@@ -112,14 +112,14 @@ def peek_image_info(img_file_names: List[str]) -> Tuple[int, int]:
     fits_sizes = list(
         map(
             get_fits_image_size,
-            filter(lambda f: f.endswith("fits"), img_file_names),
+            filter(lambda f: (f.endswith("fits") or f.endswith("fits.gz")), img_file_names),
         )
     )
 
     standard_sizes = list(
         map(
             get_standard_image_size,
-            filterfalse(lambda f: f.endswith("fits"), img_file_names),
+            filterfalse(lambda f: (f.endswith("fits") or f.endswith("fits.gz")), img_file_names),
         )
     )