-
Notifications
You must be signed in to change notification settings - Fork 0
/
skin_color_segmentation.py
145 lines (104 loc) · 4.95 KB
/
skin_color_segmentation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__ = 'Will Brennan'
import logging
import cv2
import numpy
import matplotlib.pyplot as plt
from multiprocessing import Pool
logger = logging.getLogger('main')
def get_hsv_mask(img, debug=False):
assert isinstance(img, numpy.ndarray), 'image must be a numpy array'
assert img.ndim == 3, 'skin detection can only work on color images'
logger.debug('getting hsv mask')
lower_thresh = numpy.array([0, 50, 0], dtype=numpy.uint8)
upper_thresh = numpy.array([120, 150, 255], dtype=numpy.uint8)
img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
msk_hsv = cv2.inRange(img_hsv, lower_thresh, upper_thresh)
msk_hsv[msk_hsv < 128] = 0
msk_hsv[msk_hsv >= 128] = 1
return msk_hsv.astype(float)
def get_rgb_mask(img, debug=False):
assert isinstance(img, numpy.ndarray), 'image must be a numpy array'
assert img.ndim == 3, 'skin detection can only work on color images'
logger.debug('getting rgb mask')
lower_thresh = numpy.array([45, 52, 108], dtype=numpy.uint8)
upper_thresh = numpy.array([255, 255, 255], dtype=numpy.uint8)
mask_a = cv2.inRange(img, lower_thresh, upper_thresh)
mask_b = 255 * ((img[:, :, 2] - img[:, :, 1]) / 20)
mask_c = 255 * ((numpy.max(img, axis=2) - numpy.min(img, axis=2)) / 20)
# msk_rgb = cv2.bitwise_and(mask_c, cv2.bitwise_and(mask_a, mask_b))
mask_d = numpy.bitwise_and(numpy.uint64(mask_a), numpy.uint64(mask_b))
msk_rgb = numpy.bitwise_and(numpy.uint64(mask_c), numpy.uint64(mask_d))
msk_rgb[msk_rgb < 128] = 0
msk_rgb[msk_rgb >= 128] = 1
return msk_rgb.astype(float)
def get_ycrcb_mask(img, debug=False):
assert isinstance(img, numpy.ndarray), 'image must be a numpy array'
assert img.ndim == 3, 'skin detection can only work on color images'
# logger.debug('getting ycrcb mask')
lower_thresh = numpy.array([90, 100, 130], dtype=numpy.uint8)
upper_thresh = numpy.array([230, 120, 180], dtype=numpy.uint8)
img_ycrcb = cv2.cvtColor(img, cv2.COLOR_RGB2YCR_CB)
msk_ycrcb = cv2.inRange(img_ycrcb, lower_thresh, upper_thresh)
msk_ycrcb[msk_ycrcb < 128] = 0
msk_ycrcb[msk_ycrcb >= 128] = 1
return msk_ycrcb.astype(float)
def grab_cut_mask(img_col, mask, debug=False):
assert isinstance(img_col, numpy.ndarray), 'image must be a numpy array'
assert isinstance(mask, numpy.ndarray), 'mask must be a numpy array'
assert img_col.ndim == 3, 'skin detection can only work on color images'
assert mask.ndim == 2, 'mask must be 2D'
kernel = numpy.ones((50, 50), numpy.float32) / (50 * 50)
dst = cv2.filter2D(mask, -1, kernel)
dst[dst != 0] = 255
free = numpy.array(cv2.bitwise_not(dst), dtype=numpy.uint8)
grab_mask = numpy.zeros(mask.shape, dtype=numpy.uint8)
grab_mask[:, :] = 2
grab_mask[mask == 255] = 1
grab_mask[free == 255] = 0
if numpy.unique(grab_mask).tolist() == [0, 1]:
logger.debug('conducting grabcut')
bgdModel = numpy.zeros((1, 65), numpy.float64)
fgdModel = numpy.zeros((1, 65), numpy.float64)
if img_col.size != 0:
mask, bgdModel, fgdModel = cv2.grabCut(img_col, grab_mask, None, bgdModel, fgdModel, 5,
cv2.GC_INIT_WITH_MASK)
mask = numpy.where((mask == 2) | (mask == 0), 0, 1).astype(numpy.uint8)
else:
logger.warning('img_col is empty')
return mask
def closing(mask):
assert isinstance(mask, numpy.ndarray), 'mask must be a numpy array'
assert mask.ndim == 2, 'mask must be a greyscale image'
logger.debug("closing mask of shape {0}".format(mask.shape))
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2)
return mask
def process(img, thresh=0.5, debug=False):
assert isinstance(img, numpy.ndarray), 'image must be a numpy array'
assert img.ndim == 3, 'skin detection can only work on color images'
# logger.debug("processing image of shape {0}".format(img.shape))
img = cv2.GaussianBlur(img,(3,3),1.0)
# resize image for faster computation
img = cv2.resize(img, (512, 256))
mask_hsv = get_hsv_mask(img, debug=debug)
mask_rgb = get_rgb_mask(img, debug=debug)
mask_ycrcb = get_ycrcb_mask(img, debug=debug)
n_masks = 3.0
mask = (mask_hsv + mask_rgb + mask_ycrcb ) / n_masks
mask[mask < thresh] = 0.0
mask[mask >= thresh] = 255.0
mask = mask.astype(numpy.uint8)
mask = closing(mask)
mask = grab_cut_mask(img, mask, debug=debug)
mask[mask > 0] = 1
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
img[mask == 0] = 0
# nonzero_indices = numpy.nonzero(img)
# min_y, min_x = numpy.min(nonzero_indices , axis=1)
# max_y, max_x = numpy.max(nonzero_indices , axis=1)
# img =img[min_y:max_y+1, min_x:max_x+1]
return img