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