python depth generator: no numba, imageio -> PIL, requirements.txt

This commit is contained in:
Rostislav Vasilikhin 2022-06-26 22:33:31 +02:00
parent bf8f7b4e57
commit d56e8f053f
2 changed files with 19 additions and 19 deletions

View File

@ -0,0 +1,3 @@
numpy>=1.17.4
scipy>=1.3.3
PIL>=7.0.0

View File

@ -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))