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 @@
+
@@ -37,7 +38,7 @@
-
+
@@ -92,26 +93,30 @@
-
+
-
-
-
+
+
+
-
-
-
+
+
+
+
-
-
-
+
+
+
+
-
-
-
+
+
+
+
-
+
+
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__':