Search code examples
pythonarraysloopsnumpylidar

increase speed for looping through numpy array


I am attempting to segment a LiDAR point cloud after the ground has been classified. I am using numpy to create an "image" of the the point cloud (pc) and am looping through the numpy array. I would like to speed up the loop or avoid it all together. I am going to use image segmentation techniques, but first I need to run this code to create an "image" and this is the part that is taking a while. Is there a way to increase the speed of this loop or a way to avoid it?


import numpy as np
from math import ceil, floor


'''In this case:
pc = point cloud (X,Y,Z values)'''

# point cloud is in the numpy array, pc
minx,maxx,miny,maxy = floor(np.min(pc[:,0]-1)),ceil(np.max(pc[:,0]+1)),floor(np.min(pc[:,1]-1)),ceil(np.max(pc[:,1]+1))# x,y bounding box

# grid x and y direction (resolution: 0.2 meters)
gridx = np.linspace(minx,maxx,int((maxx - minx+0.2)*5),endpoint=True) 
gridy = np.linspace(miny,maxy,int((maxy - miny +0.2)*5),endpoint=True)

#shape of the new image with 0.2 meter resolution.
imgx,imgy = int((maxx-minx+0.2)*5),int((maxy - miny +0.2)*5)

# this is what will be created at the end.  It will be a binary image.
img = np.zeros((imgx,imgy))

#loop through array to generate image (this is the part that takes a while)
for x,i in enumerate(gridx):
    for y,j in enumerate(gridy):

# Test if there any points in this "grid"
        input_point = pc[np.where(((pc[:,0]>i) & (pc[:,0]<i+1))& ((pc[:,1]>j) & (pc[:,1]<j+1)))]
# if there are points, give pixel value 1.
        if input_point.shape[0]!=0:
            img[x,y]=1

print('Image made')


Thank you.


Solution

  • Here is a vectorized version that produces the same output on a random test set:

    import numpy as np
    from math import ceil, floor
    import time
    
    width = 0.2
    
    t = [time.time()]
    
    pc = np.random.uniform(-10, 10, (100, 3))
    
    # point cloud is in the numpy array, pc
    minx,maxx,miny,maxy = floor(np.min(pc[:,0]-1)),ceil(np.max(pc[:,0]+1)),floor(np.min(pc[:,1]-1)),ceil(np.max(pc[:,1]+1))# x,y bounding box
    
    # grid x and y direction (resolution: 0.2 meters)
    gridx = np.linspace(minx,maxx,int((maxx - minx+0.2)*5),endpoint=True) 
    gridy = np.linspace(miny,maxy,int((maxy - miny +0.2)*5),endpoint=True)
    
    #shape of the new image with 0.2 meter resolution.
    imgx,imgy = int((maxx-minx+0.2)*5),int((maxy - miny +0.2)*5)
    
    print('Shared ops done')
    t.append(time.time())
    
    # this is what will be created at the end.  It will be a binary image.
    img = np.zeros((imgx,imgy))
    
    #loop through array to generate image (this is the part that takes a while)
    for x,i in enumerate(gridx):
        for y,j in enumerate(gridy):
    
    # Test if there any points in this "grid"
            input_point = pc[np.where(((pc[:,0]>i) & (pc[:,0]<i+width))& ((pc[:,1]>j) & (pc[:,1]<j+width)))]
    # if there are points, give pixel value 1.
            if input_point.shape[0]!=0:
                img[x,y]=1
    
    t.append(time.time())            
    print('Image made')
    
    if width == 0.2:
        img2 = np.zeros((imgx, imgy), 'u1')
        x2, y2 = (((pc[:, :2] - (minx, miny)) * (5, 5))).astype(int).T
        img2[x2, y2] = 1
    
    elif width == 1:
        img2 = np.zeros((imgx+4, imgy+4), 'u1')
        x2, y2 = (((pc[:, :2] - (minx, miny)) * (5, 5))).astype(int).T
    
        np.lib.stride_tricks.as_strided(img2, (imgx, imgy, 5, 5), 2 * img2.strides)[x2, y2] = 1
        img2 = img2[4:, 4:]
    
    t.append(time.time())
    print('Image remade')
    
    print('took', np.diff(t), 'secs respectively')
    assert((img2==img).all())
    print('results equal')
    

    Your code produces 5x5 pixels. Is that intentional? I had to be a bit tricksy to reproduce that.

    UPDATE: added a version that makes ordinary pixels instead.

    Sample run:

    Shared ops done
    Image made
    Image remade
    took [2.29120255e-04 1.54510736e-01 1.44481659e-04] secs respectively
    results equal