提高 Numpy For 循环速度 [英] Improving Numpy For Loop Speed

问题描述

我试图找到最接近 RGB 值 (0,0,255) 的像素.我正在尝试使用 3D 毕达哥拉斯计算来计算 RGB 值中的像素与该值的距离，将它们添加到列表中，然后返回具有最低距离的值的 X 和 Y 坐标.这是我所拥有的:

I'm trying to find the pixels closest to an RGB value of (0,0,255). I'm trying to calculate the distance of the pixel in RGB values to that value using a 3D Pythagoras calculation, add them to a list, and then return the X and Y coordinates of the values that have the lowest distance. Here's what I have:

# import the necessary packages
import numpy as np
import scipy.spatial as sp
import matplotlib.pyplot as plt
import cv2
import math
from PIL import Image, ImageDraw, ImageFont

background = Image.open("test.tif").convert('RGBA')
png = background.save("test.png")

retina = cv2.imread("test.png")
#convert BGR to RGB image
retina = cv2.cvtColor(retina, cv2.COLOR_BGR2RGB)

h,w,bpp = np.shape(retina)

min1_d = float('inf')
min1_coords = (None, None)

min2_d = float('inf')
min2_coords = (None, None)


for py in range(0,h):
    for px in range (0,w):
        r = retina[py][px][0]
        g = retina[py][px][1]
        b = retina[py][px][2]
        d = math.sqrt(((r-0)**2) + ((g-0)**2) + ((255-b)**2))
        print(str(r) + "," + str(g) + "," + str(b) + ",," + str(px) + "," + str(py) + ",," + str(d))
        if d < min1_d:
            min2_d = min1_d
            min2_coords = min1_coords
            
            min1_d = d
            min1_coords = (px, py)
        elif d < min2_d: # if it's not the smallest, check if it's the second smallest
            min2_d = d
            min2_coords = (px, py)

print(min1_coords, min2_coords)

width, height = background.size
x_max = int(width)
y_max = int(height)

img = Image.new('RGBA', (x_max, y_max), (255,255,255,0))
draw = ImageDraw.Draw(img)
draw.point(min1_coords, (0,0,255))
draw.point(min2_coords, (0,0,255))

foreground = img
background.paste(foreground, (0, 0), foreground)
foreground.save("test_bluer.png")
background.save("test_bluer_composite.png")

如何加速我的 for 循环?我相信这个答案是正确的，但我不确定如何在切片时实现 px 和 py 变量 如这个答案所示.

How can I speed up my for loops? I believe this answer is on the right track, but I'm not sure how to implement the px and py variables while slicing as this answer shows.

推荐答案

你可以通过向量化 for 循环来加速你的代码:

You can speed up your code by vectorizing the for loop:

    r = retina[:,:,0]
    g = retina[:,:,1]
    b = retina[:,:,2]
    d = np.sqrt(r**2 + g**2 + (255-b)**2)

你可以找到最小值的坐标:

You can find the coordinates of the minimum with:

    min_coords = np.unravel_index(np.argmin(d), np.shape(d))

如果您想找到第二小的距离，只需将先前的最小值更改为更大的距离:

If you want to find the second smallest distance just change the previous minimum to be a larger distance:

    d[min_coords[0],min_coords[1]] = np.inf
    min_coords = np.unravel_index(np.argmin(d), np.shape(d))
    # min_coords now has the second smallest distance

这篇关于提高 Numpy For 循环速度的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！