From fd2b9a65455e3bf4f64cdfa3dd1a5bad856effe1 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Tue, 13 Oct 2020 16:12:10 +0000
Subject: [PATCH] - improved

git-svn-id: http://moon:8086/svn/software/trunk/projects/chromakey@741 b431acfa-c32f-4a4a-93f1-934dc6c82436
---
 .idea/workspace.xml | 35 +++++++++++++++-----------
 chromakey.py        | 61 +++++++++++++++++++++++++++++----------------
 2 files changed, 59 insertions(+), 37 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 1ae450f..aed2995 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -21,6 +21,7 @@
   </component>
   <component name="PropertiesComponent">
     <property name="RunOnceActivity.OpenProjectViewOnStart" value="true" />
+    <property name="last_opened_file_path" value="$PROJECT_DIR$" />
     <property name="settings.editor.selected.configurable" value="configurable.group.appearance" />
   </component>
   <component name="RunManager" selected="Python.chromakey">
@@ -37,7 +38,7 @@
       <option name="ADD_CONTENT_ROOTS" value="true" />
       <option name="ADD_SOURCE_ROOTS" value="true" />
       <option name="SCRIPT_NAME" value="$PROJECT_DIR$/chromakey.py" />
-      <option name="PARAMETERS" value="-f $PROJECT_DIR$/../../../matlab/trunk/chromakey/gallery_2.png" />
+      <option name="PARAMETERS" value="-f $PROJECT_DIR$/chromatest.png" />
       <option name="SHOW_COMMAND_LINE" value="false" />
       <option name="EMULATE_TERMINAL" value="false" />
       <option name="MODULE_MODE" value="false" />
@@ -92,26 +93,30 @@
       <screen x="64" y="26" width="1856" height="1174" />
     </state>
     <state x="621" y="294" width="732" height="644" key="#Inspections/64.26.1856.1174@64.26.1856.1174" timestamp="1602414363028" />
-    <state x="450" y="271" key="#com.intellij.execution.impl.EditConfigurationsDialog" timestamp="1602427205388">
+    <state x="161" y="194" width="1427" height="684" key="#com.intellij.execution.impl.EditConfigurationsDialog" timestamp="1602518054229">
       <screen x="64" y="26" width="1856" height="1174" />
     </state>
-    <state x="450" y="271" key="#com.intellij.execution.impl.EditConfigurationsDialog/64.26.1856.1174@64.26.1856.1174" timestamp="1602427205388" />
-    <state width="1833" height="388" key="GridCell.Tab.0.bottom" timestamp="1602436299536">
-      <screen x="64" y="26" width="1856" height="1174" />
+    <state x="161" y="194" width="1427" height="684" key="#com.intellij.execution.impl.EditConfigurationsDialog/64.26.1856.1174@64.26.1856.1174" timestamp="1602518054229" />
+    <state width="1833" height="314" key="GridCell.Tab.0.bottom" timestamp="1602598748033">
+      <screen x="64" y="27" width="1856" height="1173" />
     </state>
-    <state width="1833" height="388" key="GridCell.Tab.0.bottom/64.26.1856.1174@64.26.1856.1174" timestamp="1602436299536" />
-    <state width="1833" height="388" key="GridCell.Tab.0.center" timestamp="1602436299536">
-      <screen x="64" y="26" width="1856" height="1174" />
+    <state width="1833" height="315" key="GridCell.Tab.0.bottom/64.26.1856.1174@64.26.1856.1174" timestamp="1602518339857" />
+    <state width="1833" height="314" key="GridCell.Tab.0.bottom/64.27.1856.1173@64.27.1856.1173" timestamp="1602598748033" />
+    <state width="1833" height="314" key="GridCell.Tab.0.center" timestamp="1602598748033">
+      <screen x="64" y="27" width="1856" height="1173" />
     </state>
-    <state width="1833" height="388" key="GridCell.Tab.0.center/64.26.1856.1174@64.26.1856.1174" timestamp="1602436299536" />
-    <state width="1833" height="388" key="GridCell.Tab.0.left" timestamp="1602436299535">
-      <screen x="64" y="26" width="1856" height="1174" />
+    <state width="1833" height="315" key="GridCell.Tab.0.center/64.26.1856.1174@64.26.1856.1174" timestamp="1602518339857" />
+    <state width="1833" height="314" key="GridCell.Tab.0.center/64.27.1856.1173@64.27.1856.1173" timestamp="1602598748033" />
+    <state width="1833" height="314" key="GridCell.Tab.0.left" timestamp="1602598748032">
+      <screen x="64" y="27" width="1856" height="1173" />
     </state>
-    <state width="1833" height="388" key="GridCell.Tab.0.left/64.26.1856.1174@64.26.1856.1174" timestamp="1602436299535" />
-    <state width="1833" height="388" key="GridCell.Tab.0.right" timestamp="1602436299536">
-      <screen x="64" y="26" width="1856" height="1174" />
+    <state width="1833" height="315" key="GridCell.Tab.0.left/64.26.1856.1174@64.26.1856.1174" timestamp="1602518339857" />
+    <state width="1833" height="314" key="GridCell.Tab.0.left/64.27.1856.1173@64.27.1856.1173" timestamp="1602598748032" />
+    <state width="1833" height="314" key="GridCell.Tab.0.right" timestamp="1602598748033">
+      <screen x="64" y="27" width="1856" height="1173" />
     </state>
-    <state width="1833" height="388" key="GridCell.Tab.0.right/64.26.1856.1174@64.26.1856.1174" timestamp="1602436299536" />
+    <state width="1833" height="315" key="GridCell.Tab.0.right/64.26.1856.1174@64.26.1856.1174" timestamp="1602518339857" />
+    <state width="1833" height="314" key="GridCell.Tab.0.right/64.27.1856.1173@64.27.1856.1173" timestamp="1602598748033" />
     <state width="1833" height="388" key="GridCell.Tab.1.bottom" timestamp="1602436299537">
       <screen x="64" y="26" width="1856" height="1174" />
     </state>
diff --git a/chromakey.py b/chromakey.py
index 210fec7..d752448 100644
--- a/chromakey.py
+++ b/chromakey.py
@@ -57,10 +57,13 @@ class Chromakey:
 	def __init__(self):
 		self.matte = None
 		self.edges = None
-		self.img_final = None
+		self.norm = None
+		self.var = None
+		self.final = None
 		self.k_ref = 0.95
 		self.k_var = 1
-		self.k_gain = 1
+		self.k_norm = 1
+		self.k_edge = 1
 		self.k_cc = 1
 		self.color_key = np.array([0.9, 0.2, 0.1])
 		ap = argparse.ArgumentParser()
@@ -69,7 +72,7 @@ class Chromakey:
 		args = vars(ap.parse_args())
 		inputfilename = args["filename"]
 
-		self.img_source = np.float32(cv.imread(inputfilename)) / 255
+		self.img_source = cv.imread(inputfilename) / 255
 		if self.img_source is None:
 			print('Failed to load file:', inputfilename)
 			sys.exit(1)
@@ -79,10 +82,12 @@ class Chromakey:
 		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.createTrackbar('k_cc', 'matte', 0, 100, self.on_change_k_cc)
-		cv.setMouseCallback('source', self.on_mouse)
+		cv.createTrackbar('k_var', 'matte', 0, 5000, self.on_change_k_var)
+		cv.createTrackbar('k_norm', 'matte', 0, 200, self.on_change_k_gain)
+		cv.createTrackbar('k_edge', 'matte', 0, 200, self.on_change_k_edge)
+		cv.createTrackbar('k_cc', 'matte', 0, 200, self.on_change_k_cc)
+		cv.setMouseCallback('source', self.on_mouse_source)
+		cv.setMouseCallback('final', self.on_mouse_final)
 		self.update_controls()
 		self.update_source()
 		self.update_matte()
@@ -91,16 +96,27 @@ class Chromakey:
 		wait_key()
 		print('Done')
 
-	def on_mouse(self, event, x, y, flags, param):
+	def on_mouse_source(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 on_mouse_final(self, event, x, y, flags, param):
+		if event == cv.EVENT_LBUTTONDOWN:
+			b = self.final[:, :, 0]*255
+			g = self.final[:, :, 1]*255
+			r = self.final[:, :, 2]*255
+			a = self.matte*255
+			bgra = cv.merge((b, g, r, a))
+			cv.imwrite("image.png", bgra)
+			print("Image written")
+
 	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))
+		cv.setTrackbarPos('k_norm', 'matte', int(self.k_norm * 100))
+		cv.setTrackbarPos('k_edge', 'matte', int(self.k_edge * 100))
 		cv.setTrackbarPos('k_cc', 'matte', int(self.k_cc * 100))
 
 	def on_change_k_ref(self, v=0):
@@ -112,7 +128,11 @@ class Chromakey:
 		self.update_matte()
 
 	def on_change_k_gain(self, v=0):
-		self.k_gain = v / 100
+		self.k_norm = v / 100
+		self.update_matte()
+
+	def on_change_k_edge(self, v=0):
+		self.k_edge = v / 100
 		self.update_matte()
 
 	def on_change_k_cc(self, v=0):
@@ -124,25 +144,22 @@ 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_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.norm = np.clip(self.k_norm * np.linalg.norm(img_filt, axis=2), 0, 1)
+		self.var = np.clip(self.k_var * np.var(img_filt, axis=2), 0, 1)
 		self.update_edges()
+		self.matte = np.clip(self.var + self.edges, 0, 1)
+		self.win_matte.update(self.matte)
+		self.update_final()
 
 	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.edges = self.matte * (1 - np.exp(-1 * (1 - self.matte)))
+		self.edges = self.k_edge * self.norm * (1-self.var)
 		self.win_edges.update(self.edges)
-		self.update_final()
 
 	def update_final(self):
 		# Multiply the foreground with the alpha 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)
+		intermediate = blend(self.img_source, self.matte)
+		self.final = despill(intermediate, self.color_key, self.k_cc * self.edges)
+		self.win_final.update(self.final)
 
 
 if __name__ == '__main__':