From d56e8f053f14b979bda20245b494aae815dd79af Mon Sep 17 00:00:00 2001
From: Rostislav Vasilikhin <savuor@gmail.com>
Date: Sun, 26 Jun 2022 22:33:31 +0200
Subject: [PATCH] python depth generator: no numba, imageio -> PIL,
 requirements.txt

---
 modules/3d/misc/python/requirements.txt |  3 +++
 modules/3d/misc/python/warp_test.py     | 35 +++++++++++--------------
 2 files changed, 19 insertions(+), 19 deletions(-)
 create mode 100644 modules/3d/misc/python/requirements.txt

diff --git a/modules/3d/misc/python/requirements.txt b/modules/3d/misc/python/requirements.txt
new file mode 100644
index 0000000000..a4567695c0
--- /dev/null
+++ b/modules/3d/misc/python/requirements.txt
@@ -0,0 +1,3 @@
+numpy>=1.17.4
+scipy>=1.3.3
+PIL>=7.0.0
diff --git a/modules/3d/misc/python/warp_test.py b/modules/3d/misc/python/warp_test.py
index 40fc844506..633fa55136 100644
--- a/modules/3d/misc/python/warp_test.py
+++ b/modules/3d/misc/python/warp_test.py
@@ -1,8 +1,7 @@
 import numpy as np
 from scipy.spatial.transform import Rotation
-import imageio
-# optional, works slower w/o it
-from numba import jit
+from PIL import Image
+import os
 
 depthFactor = 5000
 psize = (640, 480)
@@ -13,16 +12,16 @@ cy = psize[1]/2-0.5
 K = np.array([[fx,  0, cx],
               [ 0, fy, cy],
               [ 0,  0,  1]])
+
 # some random transform
 rmat = Rotation.from_rotvec(np.array([0.1, 0.2, 0.3])).as_dcm()
 tmat = np.array([[-0.04, 0.05, 0.6]]).T
 rtmat = np.vstack((np.hstack((rmat, tmat)), np.array([[0, 0, 0, 1]])))
 
-#TODO: warp rgb image as well
-testDataPath = "/path/to/sources/opencv_extra/testdata"
-srcDepth = imageio.imread(testDataPath + "/cv/rgbd/depth.png")
+testDataPath = os.getenv("OPENCV_TEST_DATA_PATH", default=None)
+srcDepth = np.asarray(Image.open(testDataPath + "/cv/rgbd/depth.png"))
+srcRgb   = np.asarray(Image.open(testDataPath + "/cv/rgbd/rgb.png"))
 
-@jit
 def reproject(image, df, K):
     Kinv = np.linalg.inv(K)
     xsz, ysz = image.shape[1], image.shape[0]
@@ -40,7 +39,6 @@ def reproject(image, df, K):
             reprojected[y, x, :] = v[:]
     return reprojected
 
-@jit
 def reprojectRtProject(image, K, depthFactor, rmat, tmat):
     Kinv = np.linalg.inv(K)
     xsz, ysz = image.shape[1], image.shape[0]
@@ -55,7 +53,6 @@ def reprojectRtProject(image, K, depthFactor, rmat, tmat):
     
     return projected
 
-@jit
 def reprojectRt(image, K, depthFactor, rmat, tmat):
     Kinv = np.linalg.inv(K)
     xsz, ysz = image.shape[1], image.shape[0]
@@ -71,10 +68,10 @@ def reprojectRt(image, K, depthFactor, rmat, tmat):
     return rotated
 
 # put projected points on a depth map
-@jit
-def splat(projected, maxv):
+def splat(projected, maxv, rgb):
     xsz, ysz = projected.shape[1], projected.shape[0]
     depth = np.full((ysz, xsz), maxv, np.float32)
+    colors = np.full((ysz, xsz, 3), 0, np.uint8)
     for y in range(ysz):
         for x in range(xsz):
             p = projected[y, x, :]
@@ -84,18 +81,18 @@ def splat(projected, maxv):
                 okuv = (u >= 0 and v >= 0 and u < xsz and v < ysz)
                 if okuv and depth[v, u] > z:
                     depth[v, u] = z
-    return depth
+                    colors[v, u, :] = rgb[y, x, :]
+    return depth, colors
 
 maxv = depthFactor
-im2 = splat(reprojectRtProject(srcDepth, K, depthFactor, rmat, tmat), maxv)
-im2[im2 >= maxv] = 0
-im2 = im2*depthFactor
+dstDepth, dstRgb = splat(reprojectRtProject(srcDepth, K, depthFactor, rmat, tmat), maxv, srcRgb)
+dstDepth[dstDepth >= maxv] = 0
+dstDepth = (dstDepth*depthFactor).astype(np.uint16)
 
-imageio.imwrite(testDataPath + "/cv/rgbd/warped.png", im2)
+Image.fromarray(dstDepth).save(testDataPath + "/cv/rgbd/warpedDepth.png")
+Image.fromarray(dstRgb  ).save(testDataPath + "/cv/rgbd/warpedRgb.png")
 
 # debug
-
-outObj = False
 def outFile(path, ptsimg):
     f = open(path, "w")
     for y in range(ptsimg.shape[0]):
@@ -105,9 +102,9 @@ def outFile(path, ptsimg):
                 f.write(f"v {v[0]} {v[1]} {v[2]}\n")
     f.close()
 
+outObj = False
 if outObj:
     objdir = "/path/to/objdir/"
     outFile(objdir + "reproj_rot_proj.obj", reproject(im2, depthFactor, K))
     outFile(objdir + "rotated.obj", reprojectRt(srcDepth, K, depthFactor, rmat, tmat))
 
-