From f7f2d9aaa5e636d2ddbb086d58a313cc514acb0b Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Sat, 17 Aug 2024 15:43:10 -0700 Subject: [PATCH] Choose key label from box instead of point --- chromakey.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/chromakey.py b/chromakey.py index 185cbc9..d3e8501 100644 --- a/chromakey.py +++ b/chromakey.py @@ -6,11 +6,12 @@ from wand.image import Image import numpy as np -def chroma_key(img, num_object_colors=1, key_coords=(0,0)): +def chroma_key(img, num_object_colors=1, key_box=((0,0),(100,100))): hsv = _hsv_points(img) hsv_cart = _hsv_to_cartesian(hsv) labels = _cluster(hsv_cart, num_object_colors+1, img) - mask = _create_mask(labels, key_coords) + key_label = _choose_key(labels, key_box) + mask = _create_mask(labels, key_label) _apply_mask(img, mask) @@ -38,13 +39,18 @@ def _cluster(points, num_clusters, img): return kmeans.labels_.reshape(img.height, img.width) +def _choose_key(labels, key_box): + box_labels = labels[key_box[0][0]:key_box[1][0], key_box[0][1]:key_box[1][1]] + values, counts = np.unique(box_labels.flatten(), return_counts=True) + i = np.argmax(counts) + return values[i] + + def _islands(mask): z = np.zeros_like(mask) - -def _create_mask(labels, key_coords): - key_label = labels[key_coords] +def _create_mask(labels, key_label): return (labels != key_label).astype(np.uint8) * 255