Numpy point cloud to image

Question

I have a point cloud which looks something like this:

The red dots are the points, the black dots are the red dots projected to the xy plane. Although it is not visible in the plot, each point also has a value, which is added to the given pixel when the point is moved to the xy plane. The points are represented by a numpy (np) array like so:

points=np.array([[x0,y0,z0,v0],[x1,y1,z1,v1],...[xn,yn,zn,vn]])

The obvious way to put these points into some image would be through a simple loop, like so:

image=np.zeros(img_size)

for point in points:
    #each point = [x,y,z,v]
    image[tuple(point[0:2])] += point[3]

Now this works fine, but it is very slow. So I was wondering if there is some way using vectorization, slicing and other clever numpy/python tricks of speeding it up, since in reality I have to this many times for large point clouds. I had come up with something using np.put:

def points_to_image(xs, ys, vs, img_size):
    img = np.zeros(img_size)
    coords = np.stack((ys, xs))
    #put the 2D coordinates into linear array coordinates
    abs_coords = np.ravel_multi_index(coords, img_size)
    np.put(img, abs_coords, ps)
    return img

(in this case the points are pre-split into vectors containing the x, y and v components). While this works fine, it of course only puts the last point to each given pixel, i.e. it is not additive.

Many thanks for your help!

user_6095728 · Answer

Courtesy of @Paul Panzer:

def points_to_image(xs, ys, ps, img_size):
    coords = np.stack((ys, xs))
    abs_coords = np.ravel_multi_index(coords, img_size)
    img = np.bincount(abs_coords, weights=ps, minlength=img_size[0]*img_size[1])
    img = img.reshape(img_size)

On my machine, the loop version takes 0.4432s vs 0.0368s using vectorization. So a neat 12x speedup.

============ EDIT ============

Quick update: using torch...

def points_to_image_torch(xs, ys, ps, sensor_size=(180, 240)):
    xt, yt, pt = torch.from_numpy(xs), torch.from_numpy(ys), torch.from_numpy(ps)
    img = torch.zeros(sensor_size)
    img.index_put_((yt, xt), pt, accumulate=True)
    return img

I get all the way down to 0.00749. And that's still all happening on CPU, so 59x speedup vs python loop. I also had a go at running it on GPU, it doesn't seem to make a difference in speed, I guess with accumulate=True it's probably using some sort of atomics on the GPU that slows it all down.

Numpy point cloud to image

1 Answers