allow images with more than 4 channels in get_image_size and read_image

opendm/ai.py

@@ -6,14 +6,21 @@
 import sys
 import rawpy
 import cv2
+from pathlib import Path
+import rasterio as rio
 
 def read_image(img_path):
-    if img_path[-4:].lower() in [".dng", ".raw", ".nef"]:
+    extension = Path(img_path).suffix.lower()
+    if extension in [".dng", ".raw", ".nef"]:
         try:
             with rawpy.imread(img_path) as r:
                 img = r.postprocess(output_bps=8, use_camera_wb=True, use_auto_wb=False)
-        except:
+        except Exception:
             return None
+    elif extension == ".tif":
+        with rio.open(img_path) as f:
+            # rasterio has the channel count as the first dimension
+            img = f.read().transpose(1,2,0)
     else:
         img = cv2.imread(img_path, cv2.IMREAD_COLOR)
         if img is None:

opendm/get_image_size.py

@@ -3,6 +3,8 @@
 import cv2
 import rawpy
 from opendm import log
+from pathlib import Path
+import rasterio as rio
 
 Image.MAX_IMAGE_PIXELS = None
 
@@ -12,10 +14,14 @@ def get_image_size(file_path, fallback_on_error=True):
     """
 
     try:
-        if file_path[-4:].lower() in [".dng", ".raw", ".nef"]:
+        extension = Path(file_path).suffix.lower()
+        if extension in [".dng", ".raw", ".nef"]:
             with rawpy.imread(file_path) as img:
                 s = img.sizes
                 width, height = s.raw_width, s.raw_height
+        elif extension == ".tif":
+            with rio.open(file_path) as f:
+                width, height = f.width, f.height
         else:
             with Image.open(file_path) as img:
                 width, height = img.size