jens
5ba397f6d6
git-svn-id: http://moon:8086/svn/software/trunk/projects/chromakey@738 b431acfa-c32f-4a4a-93f1-934dc6c82436
130 lines
3.3 KiB
Python
130 lines
3.3 KiB
Python
#!/usr/bin/env python
|
|
'''
|
|
Chroma keying in OpenCV
|
|
|
|
'''
|
|
# Python 2/3 compatibility
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
import sys
|
|
import argparse
|
|
import cv2 as cv
|
|
|
|
|
|
def colordistance(color1, color2):
|
|
d0 = (color1[0] - color2[0])
|
|
d1 = (color1[1] - color2[1])
|
|
d2 = (color1[2] - color2[2])
|
|
return np.math.sqrt(d0 * d0 + d1 * d1 + d2 * d2)
|
|
|
|
def blend(src, matte):
|
|
result = np.copy(src)
|
|
|
|
result[:, :, 0] = src[:, :, 0] * matte
|
|
result[:, :, 1] = src[:, :, 1] * matte
|
|
result[:, :, 2] = src[:, :, 2] * matte
|
|
|
|
return result
|
|
|
|
def wait_key():
|
|
while True:
|
|
ch = cv.waitKey()
|
|
if ch == 27:
|
|
break
|
|
|
|
class ImageWindow():
|
|
def __init__(self, name):
|
|
self.name = name
|
|
cv.namedWindow(self.name)
|
|
|
|
def update(self, img):
|
|
cv.imshow(self.name, img)
|
|
|
|
class Chromakey():
|
|
def __init__(self):
|
|
self.matte = None
|
|
self.k_ref = 0.95
|
|
self.k_var = 1
|
|
self.k_gain = 1
|
|
self.color_key = np.array([0.9, 0.2, 0.1])
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("-f", "--filename", type=str, default=None,
|
|
help="Input image")
|
|
|
|
args = vars(ap.parse_args())
|
|
inputfilename = args["filename"]
|
|
|
|
self.img_source = np.float32(cv.imread(inputfilename)) / 255
|
|
if self.img_source is None:
|
|
print('Failed to load file:', inputfilename)
|
|
sys.exit(1)
|
|
|
|
self.win_source = ImageWindow('source')
|
|
self.win_matte = ImageWindow('matte')
|
|
self.win_final = ImageWindow('final')
|
|
self.win_edges = ImageWindow('edges')
|
|
cv.createTrackbar('k_ref', 'matte', 0, 150, self.on_change_k_ref)
|
|
cv.createTrackbar('k_var', 'matte', 0, 1000, self.on_change_k_var)
|
|
cv.createTrackbar('k_gain', 'matte', 0, 2000, self.on_change_k_gain)
|
|
cv.setMouseCallback('source', self.onMouse)
|
|
self.update_controls()
|
|
self.update_source()
|
|
self.update_matte()
|
|
self.update_edges()
|
|
self.update_final()
|
|
wait_key()
|
|
print('Done')
|
|
|
|
def onMouse(self, event, x, y, flags, param):
|
|
if (event == cv.EVENT_LBUTTONDOWN):
|
|
self.color_key = self.img_source[y,x]
|
|
print("color_key", self.color_key)
|
|
self.update_matte()
|
|
|
|
def update_controls(self):
|
|
cv.setTrackbarPos('k_ref', 'matte', int(self.k_ref*100))
|
|
cv.setTrackbarPos('k_var', 'matte', int(self.k_var*100))
|
|
cv.setTrackbarPos('k_gain', 'matte', int(self.k_gain*100))
|
|
|
|
def on_change_k_ref(self, v=0):
|
|
self.k_ref = v/100
|
|
self.update_matte()
|
|
|
|
def on_change_k_var(self, v=0):
|
|
self.k_var = v/100
|
|
self.update_matte()
|
|
|
|
def on_change_k_gain(self, v=0):
|
|
self.k_gain = v/100
|
|
self.update_matte()
|
|
|
|
def update_source(self):
|
|
self.win_source.update(self.img_source)
|
|
|
|
def update_matte(self):
|
|
img_filt = self.img_source - self.k_ref * self.color_key
|
|
img_norm = np.linalg.norm(img_filt, axis=2)
|
|
img_var = np.var(img_filt, axis=2)
|
|
img_exp = self.k_gain*(1-np.exp(-self.k_var*img_var))
|
|
self.matte = np.clip(img_exp, 0, 1)
|
|
self.win_matte.update(self.matte)
|
|
self.update_edges()
|
|
|
|
def update_edges(self):
|
|
img_filt = self.img_source - self.k_ref * self.color_key
|
|
img_norm = np.linalg.norm(img_filt, axis=2)
|
|
edges = np.clip(img_norm-self.matte, 0, 1)
|
|
self.win_edges.update(edges)
|
|
self.update_final()
|
|
|
|
def update_final(self):
|
|
# Multiply the foreground with the alpha matte
|
|
self.img_final = blend(self.img_source, self.matte)
|
|
self.win_final.update(self.img_final)
|
|
|
|
if __name__ == '__main__':
|
|
print(__doc__)
|
|
ck = Chromakey()
|
|
cv.destroyAllWindows()
|