Numpy - More efficient code to calculate metric

Viewed 80

I am trying to implement this metric enter image description here

I already managed to calculate NUBN with numpy operations so that is fast, but I can't find a way to escape python slow looping to calculate the DRD part. Here is my current calculation of DRD:

def drd(im, im_gt):
    height, width = im.shape
    W = np.array([[1/math.sqrt(x**2+y**2) if x != 0 or y != 0 else 0 for x in range(-2, 3)] for y in range(-2, 3)])
    W /= W.sum()
    drd = 0
    s = []
    for y, x in zip(*np.where(im_gt != im)):
        if x > 1 and y > 1 and x + 2 < width and y + 2 < height:
            s.append(im_gt[y-2:y+3, x-2:x+3] == im_gt[y, x])
        else:
            for yy in range(y-2, y+3):
                for xx in range(x-2, x+3):
                    if xx > 1 and yy > 1 and xx < width - 1 and yy < height - 1:
                        drd += abs(im_gt[yy, xx] - im[y, x]) * W[yy-y+2, xx-x+2]
    return drd + np.sum(s * W)

drd(np.random.choice([0, 1], size=(100, 100)), np.random.choice([0, 1], size=(100, 100)))

Can anyone think of a faster way to do this? Timings on 1000x1000:

enter image description here

1 Answers

The first step in speeding things up with numpy is to break up your sequence of operations into something that can be applied to an entire array. Let's start with an easy one: removing the comprehensions in the computation of W:

W = np.hypot(np.arange(-2, 3), np.arange(-2, 3)[:, None])
np.reciprocal(W, where=W.astype(bool), out=W)
W /= W.sum()

The next thing (which is hinted at above with where=W.astype(bool)) is to use masking where appropriate to apply a condition to an entire array. Your algorithm is as follows:

For each location that does not match between im and im_gt, compute the sum of the elements of W centered on that location where they do not match.

You can compute this with a convolution with W. Locations where im == im_gt are simply discarded. Locations where im_gt == 1 need to be flipped by subtracting from W.sum(), since you need to sum the zeros, not the ones for those elements. Convolution is implemented in scipy.signal.convolve2d. You get the same edge effects by using mode='same' and adjusting the edge pixels carefully. You can cheat and get the edge sums by convolving with an array of ones:

from scipy.signal import convolve2d

# Compute this once outside the function
W = np.hypot(np.arange(-2, 3), np.arange(-2, 3)[:, None])
np.reciprocal(W, where=W.astype(bool), out=W)
W /= W.sum()

def drd(im, im_gt):
    m0 = im != im_gt
    m1 = im_gt == 0
    m2 = im_gt == 1
    s1 = convolve2d(m1, W, mode='same')[m0 & m1].sum()
    s2 = convolve2d(m2, W, mode='same')[m0 & m2].sum()
    return s1 + s2
Related