- removed warnings

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

chromakey.py

@@ -1,8 +1,6 @@
 #!/usr/bin/env python
-'''
-Chroma keying in OpenCV
 
-'''
+
 # Python 2/3 compatibility
 from __future__ import print_function
 
@@ -12,12 +10,13 @@ import argparse
 import cv2 as cv
 
 
-def colordistance(color1, color2):
+def color_distance(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)
 
@@ -27,13 +26,25 @@ def blend(src, matte):
 
 	return result
 
+
+def despill(src, color_bg, edges):
+	result = np.copy(src)
+
+	result[:, :, 0] = src[:, :, 0] - edges * color_bg[0]
+	result[:, :, 1] = src[:, :, 1] - edges * color_bg[1]
+	result[:, :, 2] = src[:, :, 2] - edges * color_bg[2]
+
+	return result
+
+
 def wait_key():
 	while True:
 		ch = cv.waitKey()
 		if ch == 27:
 			break
 
-class ImageWindow():
+
+class ImageWindow:
 	def __init__(self, name):
 		self.name = name
 		cv.namedWindow(self.name)
@@ -41,16 +52,19 @@ class ImageWindow():
 	def update(self, img):
 		cv.imshow(self.name, img)
 
-class Chromakey():
+
+class Chromakey:
 	def __init__(self):
 		self.matte = None
+		self.edges = None
+		self.img_final = None
 		self.k_ref = 0.95
 		self.k_var = 1
 		self.k_gain = 1
+		self.k_cc = 1
 		self.color_key = np.array([0.9, 0.2, 0.1])
 		ap = argparse.ArgumentParser()
-		ap.add_argument("-f", "--filename", type=str, default=None,
+		ap.add_argument("-f", "--filename", type=str, default=None, help="Input image")
-						help="Input image")
 
 		args = vars(ap.parse_args())
 		inputfilename = args["filename"]
@@ -67,7 +81,8 @@ class Chromakey():
 		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)
+		cv.createTrackbar('k_cc', 'matte', 0, 100, self.on_change_k_cc)
+		cv.setMouseCallback('source', self.on_mouse)
 		self.update_controls()
 		self.update_source()
 		self.update_matte()
@@ -76,27 +91,32 @@ class Chromakey():
 		wait_key()
 		print('Done')
 
-	def onMouse(self, event, x, y, flags, param):
+	def on_mouse(self, event, x, y, flags, param):
-		if (event == cv.EVENT_LBUTTONDOWN):
+		if event == cv.EVENT_LBUTTONDOWN:
-			self.color_key = self.img_source[y,x]
+			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_ref', 'matte', int(self.k_ref * 100))
-		cv.setTrackbarPos('k_var', 'matte', int(self.k_var*100))
+		cv.setTrackbarPos('k_var', 'matte', int(self.k_var * 100))
-		cv.setTrackbarPos('k_gain', 'matte', int(self.k_gain*100))
+		cv.setTrackbarPos('k_gain', 'matte', int(self.k_gain * 100))
+		cv.setTrackbarPos('k_cc', 'matte', int(self.k_cc * 100))
 
 	def on_change_k_ref(self, v=0):
-		self.k_ref = v/100
+		self.k_ref = v / 100
 		self.update_matte()
 
 	def on_change_k_var(self, v=0):
-		self.k_var = v/100
+		self.k_var = v / 100
 		self.update_matte()
 
 	def on_change_k_gain(self, v=0):
-		self.k_gain = v/100
+		self.k_gain = v / 100
+		self.update_matte()
+
+	def on_change_k_cc(self, v=0):
+		self.k_cc = v / 100
 		self.update_matte()
 
 	def update_source(self):
@@ -104,25 +124,27 @@ class Chromakey():
 
 	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_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))
+		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_filt = self.img_source - self.k_ref * self.color_key
-		img_norm = np.linalg.norm(img_filt, axis=2)
+		# img_norm = np.linalg.norm(img_filt, axis=2)
-		edges = np.clip(img_norm-self.matte, 0, 1)
+		# edges = np.clip(img_norm-self.matte, 0, 1)
-		self.win_edges.update(edges)
+		self.edges = self.matte * (1 - np.exp(-1 * (1 - self.matte)))
+		self.win_edges.update(self.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.img_final = despill(blend(self.img_source, self.matte), self.color_key, self.k_cc * self.edges)
 		self.win_final.update(self.img_final)
 
+
 if __name__ == '__main__':
 	print(__doc__)
 	ck = Chromakey()