Sky/bg filter support for DNGs

opendm/ai.py

@@ -4,6 +4,24 @@ from opendm import log
 import zipfile
 import time
 import sys
+import rawpy
+
+def read_image(img_path):
+    if img_path[-4:].lower() in [".dng", ".raw", ".nef"]:
+        try:
+            with rawpy.imread(img_path) as r:
+                img = r.postprocess(output_bps=8, use_auto_wb=True)
+        except:
+            return None
+    else:
+        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        if img is None:
+            return None
+
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    
+    return img
+
 
 def get_model(namespace, url, version, name = "model.onnx"):
     version = version.replace(".", "_")

opendm/bgfilter.py

@@ -5,6 +5,7 @@ import cv2
 import os
 import onnxruntime as ort
 from opendm import log
+from opendm.ai import read_image
 from threading import Lock
 
 mutex = Lock()
@@ -73,11 +74,7 @@ class BgFilter():
         return output
 
     def run_img(self, img_path, dest):
-        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-        if img is None:
-            return None
-
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = read_image(img_path)
         mask  = self.get_mask(img)
         
         img_name = os.path.basename(img_path)

opendm/skyremoval/skyfilter.py

@@ -6,6 +6,7 @@ import os
 import onnxruntime as ort
 from .guidedfilter import guided_filter
 from opendm import log
+from opendm.ai import read_image
 from threading import Lock
 
 mutex = Lock()
@@ -72,11 +73,7 @@ class SkyFilter():
 
     def run_img(self, img_path, dest):
 
-        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-        if img is None:
-            return None
-
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = read_image(img_path)
         img = np.array(img / 255., dtype=np.float32)
 
         mask  = self.get_mask(img)

stages/dataset.py

@@ -216,7 +216,7 @@ class ODMLoadDatasetStage(types.ODM_Stage):
                     # Generate list of sky images
                     sky_images = []
                     for p in photos:
-                        if p.mask is None and (args.camera_lens in ['fisheye', 'spherical'] or p.pitch is None or (abs(p.pitch) > 20)) and (not " " in p.filename) and not p.is_raw_format():
+                        if p.mask is None and (args.camera_lens in ['fisheye', 'spherical'] or p.pitch is None or (abs(p.pitch) > 20)) and (not " " in p.filename):
                             sky_images.append({'file': os.path.join(images_dir, p.filename), 'p': p})
 
                     if len(sky_images) > 0:
@@ -257,7 +257,7 @@ class ODMLoadDatasetStage(types.ODM_Stage):
                     # Generate list of sky images
                     bg_images = []
                     for p in photos:
-                        if p.mask is None and (not " " in p.filename) and not p.is_raw_format():
+                        if p.mask is None and (not " " in p.filename):
                             bg_images.append({'file': os.path.join(images_dir, p.filename), 'p': p})
 
                     if len(bg_images) > 0: