Color segmentation using ArcPy

import numpy as np
import math

def find_segments(ras_a, a, sd, k):
	# get useful information from ras_a.shape
	nbands = ras_a.shape[0]
	nrows = ras_a.shape[1]
	ncols = ras_a.shape[2]
	
	# calculate the radius of the target color sphere
	radius = 0
	for b in range(nbands):
		radius += (k*sd[b])**2
	radius = math.sqrt(radius)

	# create a new zero distance array
	dist_a = np.zeros(ras_a.shape[1:3])
	
	# clone the original raster array
	ret_a = ras_a.copy()
	
	# iterations
	for r in range(nrows):
		for c in range(ncols):
			# calculate the color distance for this cell at (r, c)
			for b in range(nbands):
				dist_a[r,c] += (ras_a[b,r,c]-a[b])**2
			
			# square root to get the distance, not sum of squared distances
			dist_a[r,c] = math.sqrt(dist_a[r,c])
			
			# if this cell is outside the sphere
			if dist_a[r,c] > radius:
				# return a negative distance
				dist_a[r,c] = -dist_a[r,c]
				
				# return black
				for b in range(nbands):
					ret_a[b,r,c] = 0
	
	return dist_a, ret_a

# get the map
ras = arcpy.Raster('color-map.tif')

# convert the raster object to a numpy array
ras_a = arcpy.RasterToNumPyArray(ras)

# average sample color per band
a = [55.4, 74.4, 62.6]

# sample standard deviation per band
sd = [4.336, 2.302, 3.050]

# this segmentation parameter controls the radius of the target color sphere
k = 1.0

dist_a, ret_a = find_segments(ras_a, a, sd, k)

dist = arcpy.NumPyArrayToRaster(dist_a, ras.extent.lowerLeft, ras.meanCellWidth, ras.meanCellHeight)
ret = arcpy.NumPyArrayToRaster(ret_a, ras.extent.lowerLeft, ras.meanCellWidth, ras.meanCellHeight)