Split out _apply_mask
This commit is contained in:
26
chromakey.py
26
chromakey.py
@@ -1,17 +1,17 @@
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
from sklearn.cluster import KMeans
|
from sklearn.cluster import KMeans
|
||||||
|
from wand.color import Color
|
||||||
from wand.image import Image
|
from wand.image import Image
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
def chroma_key(img, object_colors=1, key_coords=(0,0)):
|
def chroma_key(img, num_object_colors=1, key_coords=(0,0)):
|
||||||
hsv = _hsv_points(img)
|
hsv = _hsv_points(img)
|
||||||
hsv_cart = _hsv_to_cartesian(hsv)
|
hsv_cart = _hsv_to_cartesian(hsv)
|
||||||
labels = _cluster(hsv_cart, object_colors+1, img)
|
labels = _cluster(hsv_cart, num_object_colors+1, img)
|
||||||
mask = _mask(labels, key_coords)
|
mask = _create_mask(labels, key_coords)
|
||||||
with Image.from_array(mask) as mask_img:
|
_apply_mask(img, mask)
|
||||||
img.composite_channel('all_channels', mask_img, 'multiply')
|
|
||||||
|
|
||||||
|
|
||||||
def _hsv_points(img):
|
def _hsv_points(img):
|
||||||
@@ -38,14 +38,26 @@ def _cluster(points, num_clusters, img):
|
|||||||
return kmeans.labels_.reshape(img.height, img.width)
|
return kmeans.labels_.reshape(img.height, img.width)
|
||||||
|
|
||||||
|
|
||||||
def _mask(labels, key_coords):
|
def _islands(mask):
|
||||||
|
z = np.zeros_like(mask)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def _create_mask(labels, key_coords):
|
||||||
key_label = labels[key_coords]
|
key_label = labels[key_coords]
|
||||||
return (labels != key_label).astype(np.uint8) * 255
|
return (labels != key_label).astype(np.uint8) * 255
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_mask(img, mask):
|
||||||
|
with Image.from_array(mask) as mask_img:
|
||||||
|
img.alpha_channel = 'activate'
|
||||||
|
img.composite_channel('alpha', mask_img, 'copy_alpha', 0, 0)
|
||||||
|
img.background_color = Color('transparent')
|
||||||
|
|
||||||
|
|
||||||
path = 'data/RAW/SC1/BR1'
|
path = 'data/RAW/SC1/BR1'
|
||||||
filename = os.path.join(path, '41244/back.orf')
|
filename = os.path.join(path, '41244/back.orf')
|
||||||
|
|
||||||
with Image(filename=filename) as img:
|
with Image(filename=filename) as img:
|
||||||
chroma_key(img)
|
chroma_key(img)
|
||||||
img.save(filename='masked_image.jpg')
|
img.save(filename='chromakey.png')
|
||||||
Reference in New Issue
Block a user