如何优化用其他栅格区域平均值替换栅格 nan 值的 Python 循环

Question

I need to read 2 rasters, one a satellite image (target) and the other regions of this image (segmentation). The target image present numeric values and nans. The segmentation image are regions, in which each pixel with the same value are from the same region, eg, all pixels with value 1 are from segment 1.

Based on that, I would like to calculate the mean value for each segment that contains nans and replace the nan calues by the segment mean value. If I have a segment of 5 pixels and target image has the values (2,nan,4,4,2), the nan value must be replaced by 3.

I have writen a script that does this. However the code is very slow at the FOR loop when I am processing large images. Based on that, I would like to know how can the loop be improved in performance.

import numpy
import rasterio

### returns which pixels are nan
def get_gaps(img):
    gaps = numpy.argwhere( numpy.isnan( img ) )
    return( gaps )

def fill(img_targ, gaps_targ, img_seg1):
    ### Get which segments contains NA on target image
    indices_gap_targ = numpy.array( list( zip(gaps_targ[:,0], gaps_targ[:,1] ) ) )
    segments_targ = img_seg1[ indices_gap_targ[:,0], indices_gap_targ[:,1] ]
    segments_targ = numpy.unique( segments_targ[~numpy.isnan( segments_targ )] )

    for seg in segments_targ:
        ### Get seg pixel position
        seg_pixels = numpy.nonzero( img_seg1 == seg )
        seg_indices = numpy.array( list( zip(seg_pixels[:][0], seg_pixels[:][1] ) ) )
        ### Get targ pix values
        targ_values_seg = img_targ[ seg_indices[:,0], seg_indices[:,1] ]
        ### Check if any is not nan otherwise it will not have any value to use as mean
        if( numpy.any( ~numpy.isnan(targ_values_seg) ) ):
            ### Get nan position and replace by mean value
            nan_pos = numpy.isnan( targ_values_seg )
            img_targ[ seg_indices[:,0][nan_pos], seg_indices[:,1][nan_pos] ] = numpy.nanmean(targ_values_seg)
    return img_targ

input_targ_filename = "/home/path/target.tif"
input_seg1_filename = "/home/path/segmentation.tif"

with rasterio.open(input_targ_filename) as dataset:
    img_targ = dataset.read(1)
    img_targ[ img_targ < -100000 ] = numpy.nan
    kwargs = dataset.meta

with rasterio.open(input_seg1_filename) as dataset:
    img_seg1 = dataset.read(1)
    img_seg1[ img_seg1 < -100000 ] = numpy.nan

gaps_targ = get_gaps(img_targ)
img_filled = fill(img_targ, gaps_targ, img_seg1)

Answer 1

np.bincount is the tool of choice for this kind of problem. (It does essentially the same as the more intuitive np.add.at but is typically way faster.)

import numpy as np

# create mock data (this takes longer than the actual processing)
print("creating example")
N = 1000
NS = 2000
tgt = np.random.randn(N,N)
tgt[np.random.random((N,N))<0.1] = np.nan
seg = np.zeros((N,N),int)
seg.ravel()[np.random.choice(N*N,NS,replace=False)] = np.arange(1,NS+1)
idcs = np.s_[1:],np.s_[:,1:],np.s_[:-1],np.s_[:,:-1]
while np.count_nonzero(seg) < N*N/2:
    i = np.random.randint(4)
    idx,cidx = idcs[i],idcs[i-2]
    seg[idx][seg[idx]==0] = seg[cidx][seg[idx]==0]

# replace nans (in-place, overwrites nans in tgt)
print("replacing nans")
n = np.isnan(tgt)
nn = ~n
segnn = seg[nn]
tgt[n] = (np.bincount(segnn,tgt[nn],NS+1)/np.bincount(segnn,None,NS+1))[seg[n]]

# check
print("verifying",end=" ... ")
sample = np.random.randint(0,NS+1,10)
for i in sample:
    assert np.allclose(tgt[n][seg[n]==i],np.mean(tgt[nn][seg[nn]==i]))
print("looks ok")

Answer 2

to answer your example, you can do it like this with no loops:

seg = np.array((2,np.nan,4,4,2))
seg[np.isnan(seg)] = np.nanmean(seg)

output:

array([2., 3., 4., 4., 2.])

I hope this principle helps you implement this into your larger code