If you convert your image into an array, you can access the RGB values for one pixel, or one of the R, G, or B values for all pixels:
from __future__ import division
import numpy as np
from PIL import Image
im = Image.open(imfile)
arr = np.asarray(im)
arr[..., 0] # All Red values
arr[..., 1] # All Green values
arr[..., 2] # All Blue values
arr[0, 0] # RGB for first corner pixel
arr[m, n] # RGB for pixel at [m, n]
arr[m, n, 0] # R value for pixel [m, n]
arr[m, n, c] # value for color c at pixel [m, n]
You can get a ranking for each pixel using argsort
, as in:
def transform(a, c=255/2):
return c - np.abs(c - a)
ranking = transform(arr).argsort(axis=-1)
which ranks the criterion values from smallest to largest value along the last (color) axis. So this gives a new array where each 'color' array instead of being the RGB values are the sorting of the transformed R, B, and G values (call them "R', B', G'"), so if the corner pixel had G' > B' > R'
, then ranking[0, 0]
would be [0, 2, 1]
because R' (0
) is smallest, then next is B' (2
), finally the largest is G' (1
).
The advantage of doing the above is that you have an array that says which method to use on which pixel. It can have at most six orderings of the transformed channels. I suggest defining a separate function like so for each of orderings. Then, only one decision must be made within the function (the second nested if/else in your example), and it can be done with np.where
which applies one thing to parts of an array where a condition is met, and another thing to the rest. This only works for two options, but if there are multiple options (if/elif/else), other techniques can work equally well.
def bgr(a):
""" for when B' < G' < R'
"""
t = transform(a)
red, green, blue = a.transpose([2,0,1])
# same as: red, green, blue = a[..., 0], a[..., 1], a[..., 2]
r_max, g_max, b_max = t.transpose([2,0,1])
assert np.all((b_max <= g_max) & (g_max <= r_max)), "doesn't match rank"
condition = r_max >= g_max + b_max
new_red = np.where(condition, red + g_max + b_max, red + r_max)
new_green = np.where(condition, green - g_max, green - r_max + b_max)
new_blue = blue - b_max
return np.dstack([new_red, new_green, new_blue])
this function only works for the first if
in yours. I would make a new function for each of those six things, and fill them into a dict like so:
functions = {
(0, 1, 2) : rgb, # for R'<G'<B'
(0, 2, 1) : rbg, # for R'<B'<G'
#etc...
}
If your output has RGB values too:
out = np.empty_like(arr)
Then loop through all six rankings/functions:
for rank, func in functions.items():
mask = np.all(transform(arr).argsort(-1) == rank, -1)
out[mask] = func(arr[mask])