https://yulizi123.github.io/tutorials/python-basic/multiprocessing/5-pool/
import os
import numpy as np
import h5py as h5
from imageio import imread
from struct import pack, unpack
import math
import multiprocessing as mp
import glob
# extract pointcloud from rgb-d, then convert it into obj coordinate system(tsdf/poisson reconstructed object)
def im2col(im, psize):
n_channels = 1 if len(im.shape) == 2 else im.shape[0]
(n_channels, rows, cols) = (1,) * (3 - len(im.shape)) + im.shape
im_pad = np.zeros((n_channels, int(math.ceil(1.0 * rows / psize) * psize),
int(math.ceil(1.0 * cols / psize) * psize)))
im_pad[:, 0:rows, 0:cols] = im
final = np.zeros((im_pad.shape[1], im_pad.shape[2], n_channels, psize, psize))
for c in range(n_channels):
for x in range(psize):
for y in range(psize):
im_shift = np.vstack((im_pad[c, x:], im_pad[c, :x]))
im_shift = np.column_stack((im_shift[:, y:], im_shift[:, :y]))
final[x::psize, y::psize, c] = np.swapaxes(im_shift.reshape(int(im_pad.shape[1] / psize), psize,
int(im_pad.shape[2] / psize), psize), 1, 2)
return np.squeeze(final[0:rows - psize + 1, 0:cols - psize + 1])
def filterDiscontinuities(depthMap):
filt_size = 7
thresh = 1000
# Ensure that filter sizes are okay
assert filt_size % 2 == 1, "Can only use odd filter sizes."
# Compute discontinuities
offset = int((filt_size - 1) / 2)
patches = 1.0 * im2col(depthMap, filt_size)
mids = patches[:, :, offset, offset]
mins = np.min(patches, axis=(2, 3))
maxes = np.max(patches, axis=(2, 3))
discont = np.maximum(np.abs(mins - mids),
np.abs(maxes - mids))
mark = discont > thresh
# Account for offsets
final_mark = np.zeros((480, 640), dtype=np.uint16)
final_mark[offset:offset + mark.shape[0], offset:offset + mark.shape[1]] = mark
return depthMap * (1 - final_mark)
def registerDepthMap(unregisteredDepthMap, rgbImage, depthK, rgbK, H_RGBFromDepth):
unregisteredHeight = unregisteredDepthMap.shape[0]
unregisteredWidth = unregisteredDepthMap.shape[1]
registeredHeight = rgbImage.shape[0]
registeredWidth = rgbImage.shape[1]
registeredDepthMap = np.zeros((registeredHeight, registeredWidth))
xyzDepth = np.empty((4, 1))
xyzRGB = np.empty((4, 1))
# Ensure that the last value is 1 (homogeneous coordinates)
xyzDepth[3] = 1
invDepthFx = 1.0 / depthK[0, 0]
invDepthFy = 1.0 / depthK[1, 1]
depthCx = depthK[0, 2]
dept = depthK[1, 2]
rgbFx = rgbK[0, 0]
rgbFy = rgbK[1, 1]
rgbCx = rgbK[0, 2]
rgbCy = rgbK[1, 2]
undistorted = np.empty(2)
for v in range(unregisteredHeight):
for u in range(unregisteredWidth):
depth = unregisteredDepthMap[v, u]
if depth == 0:
continue
xyzDepth[0] = ((u - depthCx) * depth) * invDepthFx
d.shape[0]):
for j in range(cloud.shape[1]):
if color:
lines.append("%s %s %s %d %d %d" % tuple(cloud[i, j, :].tolist()))
else:
lines.append("%s %s %s" % tuple(cloud[i, j, :].tolist()))
for face in faces:
lines.append(("%d" + " %d" * len(face)) % tuple([len(face)] + list(face)))
f.write("\n".join(lines) + "\n")
f.close()
def writePCD(filename, pointCloud, ascii_type=True):
if len(pointCloud.shape) != 3:
print("Expected pointCloud to have 3 dimensions. Got %d instead" % len(pointCloud.shape))
return
with open(filename, "w") as f:
height = pointCloud.shape[0]
width = pointCloud.shape[1]
f.write("# .PCD v.7 - Point Cloud Data file format\n")
f.write("VERSION .7\n")
if pointCloud.shape[2] == 3:
f.write("FIELDS x y z\n")
f.write("SIZE 4 4 4\n")
f.write("TYPE F F F\n")
f.write("COUNT 1 1 1\n")
else:
f.write("FIELDS x y z rgb\n")
f.write("SIZE 4 4 4 4\n")
f.write("TYPE F F F F\n")
f.write("COUNT 1 1 1 1\n")
f.write("WIDTH %d\n" % width)
f.write("HEIGHT %d\n" % height)
f.write("VIEWPOINT 0 0 0 1 0 0 0\n")
f.write("POINTS %d\n" % (height * width))
if ascii_type:
f.write("DATA ascii\n")
for row in range(height):
for col in range(width):
if pointCloud.shape[2] == 3:
f.write("%f %f %f\n" % tuple(pointCloud[row, col, :]))
else:
f.write("%f %f %f" % tuple(pointCloud[row, col, :3]))
r = int(pointCloud[row, col, 3])
g = int(pointCloud[row, col, 4])
b = int(pointCloud[row, col, 5])
rgb_int = (r << 16) | (g << 8) | b
packed = pack("i", rgb_int)
rgb = unpack("f", packed)[0]
f.write(" %.12e\n" % rgb)
else:
f.write("DATA binary\n")
if pointCloud.shape[2] == 6:
# These are written as bgr because rgb is interpreted as a single
# little-endian float.
dt = np.dtype([("x", np.float32),
("y", np.float32),
("z", np.float32),
("b", np.uint8),
("g", np.uint8),
("r", np.uint8),
("I", np.uint8)])
pointCloud_tmp = np.zeros((height * width, 1), dtype=dt)
for i, k in enumerate(["x", "y", "z", "r", "g", "b"]):
pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height * width, 1))
pointCloud_tmp.tofile(f)
else:
dt = np.dtype([("x", np.float32),
("y", np.float32),
("z", np.float32),
("I", np.uint8)])
pointCloud_tmp = np.zeros((height * width, 1), dtype=dt)
for i, k in enumerate(["x", "y", "z"]):
pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height * width, 1))
pointCloud_tmp.tofile(f)
def getRGBFromDepthTransform(calibration, camera, referenceCamera):
irKey = "H_{0}_ir_from_{1}".format(camera, referenceCamera)
rgbKey = "H_{0}_from_{1}".format(camera, referenceCamera)
rgbFromRef = calibration[rgbKey][:]
irFromRef = calibration[irKey][:]
return np.dot(rgbFromRef, np.linalg.inv(irFromRef)), np.linalg.inv(rgbFromRef)
def generate(path):
path = path.split("/")
# Parameters
#print(path)
#['', 'data0', Files', 'PointCode', 'PointNet', 'PointNetGPD', 'data', 'ycb-tools', 'models', 'ycb', '003_cracker_box', 'rgbd', 'NP1_354.jpg']
contains the ycb data.
[1].split(".")[0] # Relative angle of the object w.r.t the camera
#print(viewpoint_angle)
# (angle of the turntable).
referenceCamera = "NP5" # can only be NP5
ply_fname = os.path.join(ycb_data_folder, target_object, "clouds",
"pc_" + viewpoint_camera + "_" + referenceCamera + "_" + viewpoint_angle + ".ply")
pcd_fname = os.path.join(ycb_data_folder, target_object, "clouds",
"pc_" + viewpoint_camera + "_" + referenceCamera + "_" + viewpoint_angle + ".pcd")
npy_fname = os.path.join(ycb_data_folder, .0001 # 100um to meters
H_RGBFromDepth, refFromRGB = getRGBFromDepthTransform(calibration, viewpoint_camera, referenceCamera)
unregisteredDepthMap = h5.File(depthFilename, "r")["depth"][:]
unregisteredDepthMap = filterDiscontinuities(unregisteredDepthMap) * depthScale
registeredDepthMap = registerDepthMap(unregisteredDepthMap, rgbImage, depthK, rgbK, H_RGBFromDepth)
# apply mask
registeredDepthMap[pbmImage == 255] = 0
pointCloud = registeredDepthMapToPointCloud(registeredDepthMap, rgbImage, rgbK, refFromRGB, objFromref)
writePLY(ply_fname, pointCloud)
writePCD(pcd_fname, pointCloud)
np.save(npy_fname, pointCloud[:, :, :3].reshape(-1, 3))
except:
f = open("exception.txt", "a")
f.write("/".join(path) + "\n")
f.close()
print(ycb_data_folder, target_object, viewpoint_camera, viewpoint_angle, "failed")
return
def main():
pointnetgpd_dir = os.environ["PointNetGPD_FOLDER"]
print(pointnetgpd_dir)
#/data0Files/PointCode/PointNet/PointNetGPD
#/data0/Files/PointCode/PointNet/PointNetGPD/data/ycb-tools/models/ycb
#/data0Files/PointCode/PointNet/PointNetGPDdata/ycb-tools/models/ycb/*/rgbd/*.jpg
#print(pointnetgpd_dir+"data/ycb-tools/models/ycb/*/rgbd/*.jpg")
fl = np.array(glob.glob(pointnetgpd_dir+"/data/ycb-tools/models/ycb/*/rgbd/*.jpg"))
np.random.shuffle(fl)
#print(len(fl)) 16800
cores = mp.cpu_count()#40
pool = mp.Pool(processes=20)
pool.map(generate, fl)
if __name__ == "__main__":
main()