How can I simplify the following code so it runs faster?

Question

I have a three-dimensional array containing many 2D images (frames). I want to remove the background by considering a threshold for each pixel value and copy new elements in a new 3D array. I wrote the following code lines, but it is too expensive for running. How can I speed up this code?

ss = stack #3D array (571, 1040, 1392)
T,ni,nj = ss.shape
Background_intensity = np.ones([T,ni,nj])
Intensity = np.zeros([T,ni,nj])
DeltaF_F_max = np.zeros([T,ni,nj])
for t in range(T):
    for i in range(ni):
        for j in range(nj):
            if ss[t,i,j]<12:
                Background_intensity[t,i,j] = ss[t,i,j]
                if Background_intensity[t,i,j] == 0 :
                    Background_intensity[t,i,j] = 1
            else:
                Intensity[t,i,j] = ss[t,i,j]
            DeltaF_F_max[t,i,j]=(((Intensity[t,i,j] - Background_intensity[t,i,j])))/(Background_intensity[t,i,j])

Answer 1

I had a go at this with Numpy. I am not sure what results you got, but it takes around 20s on my Mac. It is quite a memory hog even after I reduced all the sizes by a factor of 8 because you don't need an int64 to store a 1 or a number under 12 or 255.

I wonder if you need to do 571 images all in one go or whether you could do them "on-the-fly" as you acquire them rather than gathering them all in one enormous lump.

You could also consider doing this with Numba as it is very good at optimising for loops - try putting [numba] in the search box above, or looking at this example - using prange to parallelise the loops across your CPU cores.

Anyway, here is my code:

#!/usr/bin/env python3

# https://stackoverflow.com/q/71460343/2836621

import numpy as np

T, ni, nj = 571, 1040, 1392

# Create representative input data, such that around 1/3 of it is < 12 for testing
ss = np.random.randint(0,36,(T,ni,nj), np.uint8)

# Ravel into 1-D representation for simpler indexing
ss_r = ss.ravel()

# Create extra arrays but using 800MB rather than 6.3GB each, also ravelled
Background_intensity = np.ones(T*ni*nj, np.uint8)
Intensity = np.zeros(T*ni*nj, np.uint8)

# Make Boolean (True/False) mask of elements below threshold
mask = ss_r < 12
# Quick check here - print(np.count_nonzero(mask)/np.size(ss)) and check it is 0.333

# Set Background_intensity to "ss" according to mask
Background_intensity[mask] = ss_r[mask]

# Make sure no zeroes present
Background_intensity[Background_intensity==0] = 1

# This corresponds to the "else" of your original "if" statement
Intensity[~mask] = ss_r[~mask]

# Final calculation and reshaping back to original shape
DeltaF_F_max = (Intensity - Background_intensity)/Background_intensity
DeltaF_F_max.reshape((T,ni,nj))