CODEWITHSUNDEEP
pythonpython-3.xnumpy
wandadars
wandadars

Reputation: 1183

Faster implementation of pixel mapping

This method is very slow. The short and sweet is that it takes in a dictionary phase_color_labels that maps an arbitrary name to a 3 element list that corresponds to RGB values, and it maps each pixel of the input image to whatever pixel value that it is closest to in the phase_color_labels dictionary. I haven't figured out if there is a vectorized version of this that would run much faster.

The image variable is just a numpy array [H, W, Channels].

def map_pixels_to_discrete_values(image, phase_color_labels):
        """
        Takes an image with floating point pixel values and maps each pixel RGB value
        to a new value based on the closest Euclidean distance to one of the RGB sets
        in the phase_label input dictionary.
        """
        mapped_image = np.copy(image)
        for i in range(mapped_image.shape[0]):
            for j in range(mapped_image.shape[1]):
                min_distance = np.inf
                min_distance_label = None
                for phase_name, phase_color in phase_color_labels.items():
                    r = phase_color[0]
                    g = phase_color[1]
                    b = phase_color[2]
                    rgb_distance = (mapped_image[i, j, 0] - r)**2 + (mapped_image[i, j, 1] - g)**2 + (mapped_image[i, j, 2] - b)**2
                    if rgb_distance < min_distance:
                        min_distance = rgb_distance
                        min_distance_label = phase_name

                mapped_image[i, j, :] = phase_color_labels[min_distance_label]
        return mapped_image

Upvotes: 1

Views: 421

Answers (1)

jirassimok
jirassimok

Reputation: 4263

To do things fast with Numpy, you generally want to avoid loops and push as much of the work as possible into Numpy's matrix operations.

Basic idea of my answer:

  1. Get the colors from phase_color_labels as an ndarray, phase_colors.
  2. Use Numpy's broadcasting to compute the "outer distance"—the euclidean distance between each image in the array and each color in phase_colors.
  3. Find the index of the color with the lowest distance for each pixel, and use those as indices in phase_colors.
phase_colors = np.array([color for color in phase_color_labels.values()])

distances = np.sqrt(np.sum((image[:,:,np.newaxis,:] - phase_colors) ** 2, axis=3))
min_indices = distances.argmin(2)

mapped_image = phase_colors[min_indices]

The third line warrants some additional explanation. First, note that phase_names and phase_colors both have shape (L, C), where L is the number of labels and C is the number of channels.

  • image[:,:,np.newaxis,:] inserts a new axis between the second and third axes, so the resulting array has shape (H, W, 1, C).
  • When subtracting an array of shape (L, C) from an array of shape (H, W, 1, C), Numpy broadcasts the arrays to shape (H, W, L, C). You can find more details on Numpy's broadcasting semantics here.
  • Then, taking the sum along axis 3 produces an array of shape (H, W, L).
  • (Neither squaring nor taking the square root effect the shape of the array.)

In the fourth line, using argmin on axis 2 then reduces the array to shape (H, W), with each value being an index from the reduced axis L—in other words, an index into phase_colors.

As an additional improvement, because square root is a monotonically increasing function, it won't change which distance is smallest, so you can remove it entirely.

Note that with large image and phase_color_labels, the memory costs of the broadcasting may be noticeable, which may also cause performance issues.

Upvotes: 2

Related Questions