Files
chromakey/color_band.py
2025-04-17 09:28:25 +02:00

93 lines
1.9 KiB
Python

from pickletools import uint8
import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt
def key_ish(kid: int, color: np.array):
prg = [(1,2), (0,2), (0,1)]
res = 0
c = prg[kid]
ci1 = c[0]
ci2 = c[1]
c1 = color[ci1]
c2 = color[ci2]
k0_1 = color[kid]
k0_2 = color[kid]
k0_12 = color[kid]
s12 = c1 + c2
if c1 > 0:
k0_1 = color[kid] / c1
if c2 > 0:
k0_2 = color[kid] / c2
if s12 > 0:
k0_12 = color[kid] / s12
# if s12 > 2*c1 and s12 > 2*c2:
# res = k0_12
if c1 > c2:
res = k0_1
else:
res = k0_2
return res
def despill(src, key_color, mat: float):
result = np.copy(src)
result[:, :, 0] -= (1-mat)*key_color[:, :, 0]
result[:, :, 1] -= (1-mat)*key_color[:, :, 1]
result[:, :, 2] -= (1-mat)*key_color[:, :, 2]
return result
def main():
key_index = 1
key_bgr = np.float32([[[0, 0, 0]]])
key_bgr[:,:,key_index] = 1
orig = np.ones([200, 360, 3], dtype=np.float32)
for x, hn in enumerate(np.linspace(0, 1, 360)):
for y in range(200):
s = y/200
orig_hsv = np.float32([[[hn * 360, s, 1]]])
orig_bgr = cv.cvtColor(orig_hsv, cv.COLOR_HSV2BGR)
orig[y, x] = orig_bgr
cv.imshow("orig", orig)
k_tol = 0.3
k_boost = 1.5
_x = []
_y = []
pic = np.ones([200, 360, 3], dtype=np.float32)
mat = np.ones([200, 360, 1], dtype=np.float32)
for x, hn in enumerate(np.linspace(0, 1, 360)):
_x.append(x)
v = 0
for y in range(200):
s = y/200
spec_hsv = np.float32([[[hn * 360, s, 1]]])
spec_bgr = cv.cvtColor(np.array(spec_hsv, dtype=np.float32),cv.COLOR_HSV2BGR)
v0 = spec_bgr[0,0,:]
v1 = min(1.0, k_boost * (1.0 - min(1.0, k_tol*key_ish(key_index, v0))))
pic_hsv = np.float32([[[hn*360,s,v1]]])
pic_bgr = cv.cvtColor(pic_hsv,cv.COLOR_HSV2BGR)
pic_bgr = despill(pic_bgr, key_bgr, v1)
pic[y, x] = pic_bgr
mat[y, x] = v1
_y.append(v1)
print(mat)
cv.imshow("pic", pic)
cv.imshow("mat", mat)
plt.plot(_x, _y)
plt.grid()
# plt.show()
if __name__ == '__main__':
main()
ch = cv.waitKey()