added

git-svn-id: http://moon:8086/svn/software/trunk/projects/chromakey@738 b431acfa-c32f-4a4a-93f1-934dc6c82436

chromakey.py

@@ -0,0 +1,129 @@
+#!/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()