Read depth maps from OpenXR files

shapovalov · facebook-github-bot · commit d851bc317355 · 2023-07-13T04:50:03.000-07:00
Summary:
Blender uses OpenEXR to dump depth maps, so we have to support it.
OpenCV requires to explicitly accepth the vulnerabilities by setting the env var before exporting.
We can set it but I think it should be user’s responsibility.
OpenCV error reporting is adequate, so I don’t handle the error on our side.

Reviewed By: bottler

Differential Revision: D47403884

fbshipit-source-id: 2fcadd1df9d0efa0aea563bcfb2e3180b3c4d1d7
diff --git a/pytorch3d/implicitron/dataset/utils.py b/pytorch3d/implicitron/dataset/utils.py
@@ -245,10 +245,21 @@ def load_mask(path: str) -> np.ndarray:
 
 
 def load_depth(path: str, scale_adjustment: float) -> np.ndarray:
-    if not path.lower().endswith(".png"):
+    if path.lower().endswith(".exr"):
+        # NOTE: environment variable OPENCV_IO_ENABLE_OPENEXR must be set to 1
+        # You will have to accept these vulnerabilities by using OpenEXR:
+        # https://github.com/opencv/opencv/issues/21326
+        import cv2
+
+        d = cv2.imread(path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[..., 0]
+        d[d > 1e9] = 0.0
+    elif path.lower().endswith(".png"):
+        d = load_16big_png_depth(path)
+    else:
         raise ValueError('unsupported depth file name "%s"' % path)
 
-    d = load_16big_png_depth(path) * scale_adjustment
+    d = d * scale_adjustment
+
     d[~np.isfinite(d)] = 0.0
     return d[None]  # fake feature channel
 
