diff --git a/chromakey.py b/chromakey.py index d3e8501..42d71d4 100644 --- a/chromakey.py +++ b/chromakey.py @@ -6,13 +6,16 @@ from wand.image import Image import numpy as np -def chroma_key(img, num_object_colors=1, key_box=((0,0),(100,100))): +def chroma_key(img, num_object_colors=1, key_box=((0,0),(100,100)), min_object_pixels=100): hsv = _hsv_points(img) hsv_cart = _hsv_to_cartesian(hsv) labels = _cluster(hsv_cart, num_object_colors+1, img) key_label = _choose_key(labels, key_box) mask = _create_mask(labels, key_label) + ids = _islands(mask) + _filter_islands(mask, ids, min_object_pixels) _apply_mask(img, mask) + img.trim() def _hsv_points(img): @@ -47,7 +50,44 @@ def _choose_key(labels, key_box): def _islands(mask): - z = np.zeros_like(mask) + ids = np.zeros_like(mask).astype(np.uint64) + next_id = 1 + + for (i, j), _ in np.ndenumerate(mask): + _flood_fill(ids, mask, i, j, next_id) + next_id += 1 + + return ids + + +def _filter_islands(mask, ids, min_object_pixels): + values, counts = np.unique(ids.flatten(), return_counts=True) + ix = np.where(counts < min_object_pixels) + mask[np.isin(ids, values[ix])] = 0 + + +def _flood_fill(ids, mask, i, j, val): + queue = [(i, j)] + + while queue: + (i, j) = queue.pop() + + if i < 0 or i >= mask.shape[0] or j < 0 or j >= mask.shape[1]: + continue + + if mask[i, j] == 0 or ids[i, j] != 0: + continue + + ids[i, j] = val + + queue.append((i - 1, j)) + queue.append((i + 1, j)) + queue.append((i - 1, j - 1)) + queue.append((i + 1, j - 1)) + queue.append((i - 1, j + 1)) + queue.append((i + 1, j + 1)) + queue.append((i, j - 1)) + queue.append((i, j + 1)) def _create_mask(labels, key_label): @@ -65,5 +105,6 @@ path = 'data/RAW/SC1/BR1' filename = os.path.join(path, '41244/back.orf') with Image(filename=filename) as img: + img.crop(left=0, top=0, width=4000, height=3000) chroma_key(img) img.save(filename='chromakey.png') \ No newline at end of file