def solvemask(notats, tats, color):
print("prepping")
# turn image object into list of pixels
notats_data = list(notats.getdata())
tats_data = list(tats.getdata())
# convert to numpy float-mode
notats_data = np.array(notats_data, dtype="float64") / 255
tats_data = np.array(tats_data, dtype="float64") / 255
color = np.array(color) / 255
# for each pixel value, compare before & after to find what % of the specified color should be added
# do this by trying all 0-255 possible alpha values
# for each alpha value, do the alpha blending algorithm, get the resulting color, and compare against the "after" color
# component-wise R G B difference, squared, averaged to get mean-square-error
# select the alpha with the min MSE
# pre-calculate the float versions of the alpha values
alpha_list = np.arange(256, dtype="float64") / 255
alpha_list = alpha_list.reshape((256,1))
# this holds resulting alpha values
alpha = np.zeros(tats_data.shape[0], dtype="uint8")
# this is a temp buffer used to store the MSE
blend_list = np.zeros(256, dtype="float64")
mse = np.zeros(256, dtype="float64")
print("running")
for d, (before, after) in enumerate(zip(notats_data, tats_data)):
progprint(d / tats_data.shape[0])
# first check shortcuts
if np.array_equal(before, after):
alpha[d] = 0 # transparent
continue
if np.array_equal(before, color):
alpha[d] = 255 # opaque
continue
# now do the hard way
# calc resulting colors after multiplying with all possible alphas
blend_list = alpha_blend_bulk(before, color, alpha_list)
# now calculate the error
mse = blend_list - after
# then square
mse = np.square(mse)
# then mean
mse = np.mean(mse, axis=1)
# then find the lowest error
bestfit = np.argmin(mse)
alpha[d] = bestfit
# print(before, after)
progclean()
print("done iterating")
nonzero = np.count_nonzero(alpha)
print("mask covers %f%%" % (100 * nonzero / alpha.size))
opaque = np.count_nonzero(alpha == 255)
print("mask opaque %f%%" % (100 * opaque / alpha.size))
# # convert results from 1d array to 2d array
# alpha = alpha.reshape(tats.size)
return alpha
def mkcurve(chan1, chan2):
"""Calculate channel curve by averaging target values."""
sums = {}
asdf = ravel(chan1)
asdff = len(asdf)
for z, (v1, v2) in enumerate(zip(asdf, ravel(chan2))):
progprint(z / asdff)
try:
sums[v1].append(v2)
except KeyError:
sums[v1] = [v2]
c = array([(src, mean(vals)) for src, vals in sorted(sums.items())])
nvals = interp(range(256), c[:, 0], c[:, 1], 0, 255)
progclean()
return dict(zip(range(256), nvals))
def correct_bad(good, bad, read=False):
"""Match colors of the bad image to good image."""
r, g, b = bad.transpose((2, 0, 1))
r2, g2, b2 = good.transpose((2, 0, 1))
corr = bad.copy()
h, w = corr.shape[:2]
if read:
print("start r")
rc = mkcurve(r, r2)
print("start g")
gc = mkcurve(g, g2)
print("start b")
bc = mkcurve(b, b2)
print("pickling")
with open('r.curve', 'wb') as rf:
pickle.dump(rc, rf)
with open('g.curve', 'wb') as gf:
pickle.dump(gc, gf)
with open('b.curve', 'wb') as bf:
pickle.dump(bc, bf)
else:
with open('r.curve', 'rb') as rf:
rc = pickle.load(rf)
with open('g.curve', 'rb') as gf:
gc = pickle.load(gf)
with open('b.curve', 'rb') as bf:
bc = pickle.load(bf)
print("apply")
for row in range(h):
progprint(row / h)
for col in range(w):
r, g, b = corr[row, col]
corr[row, col] = [rc[r], gc[g], bc[b]]
progclean()
return corr