Choose key label from box instead of point
This commit is contained in:
16
chromakey.py
16
chromakey.py
@@ -6,11 +6,12 @@ from wand.image import Image
|
|||||||
import numpy as np
|
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 = _hsv_points(img)
|
||||||
hsv_cart = _hsv_to_cartesian(hsv)
|
hsv_cart = _hsv_to_cartesian(hsv)
|
||||||
labels = _cluster(hsv_cart, num_object_colors+1, img)
|
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)
|
_apply_mask(img, mask)
|
||||||
|
|
||||||
|
|
||||||
@@ -38,13 +39,18 @@ def _cluster(points, num_clusters, img):
|
|||||||
return kmeans.labels_.reshape(img.height, img.width)
|
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):
|
def _islands(mask):
|
||||||
z = np.zeros_like(mask)
|
z = np.zeros_like(mask)
|
||||||
|
|
||||||
|
|
||||||
|
def _create_mask(labels, key_label):
|
||||||
def _create_mask(labels, key_coords):
|
|
||||||
key_label = labels[key_coords]
|
|
||||||
return (labels != key_label).astype(np.uint8) * 255
|
return (labels != key_label).astype(np.uint8) * 255
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user