tools/human_segmenter.py

# coding=utf-8
# date:
import tensorflow as tf
import numpy as np
import cv2
import os

if tf.__version__ >= '2.0':
    print("tf version >= 2.0")
    tf = tf.compat.v1
    tf.disable_eager_execution()


class human_segmenter(object):
    def __init__(self, model_path,is_encrypted_model=False):
        super(human_segmenter, self).__init__()
        f = tf.gfile.FastGFile(model_path, 'rb')
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        persisted_graph = tf.import_graph_def(graph_def, name='')

        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = 0.3  # 占用GPU 30%的显存
        self.sess = tf.InteractiveSession(graph=persisted_graph, config=config)

        # self.image_node = self.sess.graph.get_tensor_by_name("input_image:0")
        # # self.output_node = self.sess.graph.get_tensor_by_name("output_png:0")
        # # check if the nodename in model
        # if "output_png:0" in self.sess.graph_def.node:
        #     self.output_node = self.sess.graph.get_tensor_by_name("output_png:0")
        # else:
        #     self.output_node = self.sess.graph.get_tensor_by_name("output_alpha:0")
        # if "if_person:0" in self.sess.graph_def.node:
        #     self.logits_node = self.sess.graph.get_tensor_by_name("if_person:0")

        print("human_segmenter init done")
    
    def image_preprocess(self, img):
        if len(img.shape) == 2:
            img = np.dstack((img, img, img))
        elif img.shape[2] == 4:
            img = img[:, :, :3]
        img = img[:, :, ::-1]
        img = img.astype(np.float32)
        return img
    
    def run(self, img):
        image_feed = self.image_preprocess(img)
        output_img_value, logits_value = self.sess.run([self.sess.graph.get_tensor_by_name("output_png:0"), self.sess.graph.get_tensor_by_name("if_person:0")],
                                                  feed_dict={self.sess.graph.get_tensor_by_name("input_image:0"): image_feed})
        # mask = output_img_value[:,:,-1]
        output_img_value = cv2.cvtColor(output_img_value, cv2.COLOR_RGBA2BGRA)
        return output_img_value

    def run_head(self, img):
        image_feed = self.image_preprocess(img)
        # image_feed = image_feed/255.0
        output_alpha = self.sess.run(self.sess.graph.get_tensor_by_name('output_alpha:0'),
                                     feed_dict={'input_image:0': image_feed})

        return output_alpha
    
    def get_human_bbox(self, mask):
        '''
        
        :param mask:
        :return: [x,y,w,h]
        '''
        print('dtype:{}, max:{},shape:{}'.format(mask.dtype, np.max(mask), mask.shape))
        ret, thresh = cv2.threshold(mask,127,255,0)
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 0:
            return None
        
        contoursArea = [cv2.contourArea(c) for c in contours]
        max_area_index = contoursArea.index(max(contoursArea))
        bbox = cv2.boundingRect(contours[max_area_index])
        return bbox
    
    
    def release(self):
        self.sess.close()


class head_segmenter(object):
    def __init__(self, model_path, is_encrypted_model=False):
        super(head_segmenter, self).__init__()
        f = tf.gfile.FastGFile(model_path, 'rb')
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        persisted_graph = tf.import_graph_def(graph_def, name='')

        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = 0.3  # 占用GPU 30%的显存
        self.sess = tf.InteractiveSession(graph=persisted_graph, config=config)

        print("human_segmenter init done")

    def image_preprocess(self, img):
        if len(img.shape) == 2:
            img = np.dstack((img, img, img))
        elif img.shape[2] == 4:
            img = img[:, :, :3]
        img = img[:, :, ::-1]
        img = img.astype(np.float32)
        return img

    def run_head(self, img):
        image_feed = self.image_preprocess(img)
        # image_feed = image_feed/255.0
        output_alpha = self.sess.run(self.sess.graph.get_tensor_by_name('output_alpha:0'),
                                     feed_dict={'input_image:0': image_feed})

        return output_alpha

    def get_human_bbox(self, mask):
        '''

        :param mask:
        :return: [x,y,w,h]
        '''
        print('dtype:{}, max:{},shape:{}'.format(mask.dtype, np.max(mask), mask.shape))
        ret, thresh = cv2.threshold(mask, 127, 255, 0)
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 0:
            return None

        contoursArea = [cv2.contourArea(c) for c in contours]
        max_area_index = contoursArea.index(max(contoursArea))
        bbox = cv2.boundingRect(contours[max_area_index])
        return bbox

    def release(self):
        self.sess.close()


class hair_segmenter(object):
    def __init__(self, model_dir, is_encrypted_model=False):
        head_path = os.path.join(model_dir, 'Matting_headparser_6_18.pb')
        face_path = os.path.join(model_dir, 'segment_face.pb')
        detect_path = os.path.join(model_dir, 'face_detect.pb')

        self.sess = self.load_sess(head_path)
        image = np.ones((512, 512, 3))
        output_png = self.sess.run(self.sess.graph.get_tensor_by_name('output_alpha:0'),
                                   feed_dict={'input_image:0': image})

        self.sess_detect = self.load_sess(detect_path)
        oboxes, scores, num_detections = self.sess_detect.run(
            [self.sess_detect.graph.get_tensor_by_name('tower_0/boxes:0'),
             self.sess_detect.graph.get_tensor_by_name('tower_0/scores:0'),
             self.sess_detect.graph.get_tensor_by_name('tower_0/num_detections:0')],
            feed_dict={'tower_0/images:0': image[np.newaxis], 'training_flag:0': False})
        faceRects = []

        self.sess_face = self.load_sess(face_path)
        image = np.ones((512, 512, 3))
        output_alpha = self.sess_face.run(self.sess_face.graph.get_tensor_by_name('output_alpha_face:0'),
                                          feed_dict={'input_image_face:0': image})

    def load_sess(self, model_path):
        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.allow_growth = True
        sess = tf.Session(config=config)
        with tf.gfile.FastGFile(model_path, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name='')
            sess.run(tf.global_variables_initializer())
        return sess

    def image_preprocess(self, img):
        if len(img.shape) == 2:
            img = np.dstack((img, img, img))
        elif img.shape[2] == 4:
            img = img[:, :, :3]
        img = img[:, :, ::-1]
        img = img.astype(np.float32)
        return img

    def run_head(self, image):
        image_feed = self.image_preprocess(image)
        output_img_value = self.sess.run(self.sess.graph.get_tensor_by_name('output_alpha:0'),
                                     feed_dict={'input_image:0': image_feed})
        # mask = output_img_value[:,:,-1]
        output_img_value = cv2.cvtColor(output_img_value, cv2.COLOR_RGBA2BGRA)
        return output_img_value

    def run(self, image):
        h, w, c = image.shape
        faceRects = self.detect_face(image)
        face_num = len(faceRects)
        print('face_num:{}'.format(face_num))
        all_head_alpha = []
        all_face_mask = []
        for i in range(face_num):
            y1 = faceRects[i][0]
            y2 = faceRects[i][1]
            x1 = faceRects[i][2]
            x2 = faceRects[i][3]
            pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box(y1, y2, x1, x2, 0.15, 0.15, 0.15, 0.15, h, w)
            temp_img = image.copy()
            roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2]
            output_alpha = self.sess_face.run(self.sess_face.graph.get_tensor_by_name('output_alpha_face:0'),
                                              feed_dict={'input_image_face:0': roi_img[:, :, ::-1]})
            face_mask = np.zeros((h, w, 3))
            face_mask[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha
            all_face_mask.append(face_mask)
            # cv2.imwrite(str(i)+'face.jpg',face_mask)
            # cv2.imwrite(str(i)+'face_roi.jpg',roi_img)

        for i in range(face_num):
            y1 = faceRects[i][0]
            y2 = faceRects[i][1]
            x1 = faceRects[i][2]
            x2 = faceRects[i][3]
            pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box(y1, y2, x1, x2, 1.47, 1.47, 1.3, 2.0, h, w)
            temp_img = image.copy()
            for j in range(face_num):
                y1 = faceRects[j][0]
                y2 = faceRects[j][1]
                x1 = faceRects[j][2]
                x2 = faceRects[j][3]
                small_y1, small_y2, small_x1, small_x2 = self.pad_box(y1, y2, x1, x2, -0.1, -0.1, -0.1, -0.1, h, w)
                small_width = small_x2 - small_x1
                small_height = small_y2 - small_y1
                if (
                        small_x1 < 0 or small_y1 < 0 or small_width < 3 or small_height < 3 or small_x2 > w or small_y2 > h):
                    continue
                # if(i!=j):
                #     temp_img[small_y1:small_y2,small_x1:small_x2]=0
                if (i != j):
                    temp_img = temp_img * (1.0 - all_face_mask[j] / 255.0)

            roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2]
            output_alpha = self.sess.run(self.sess.graph.get_tensor_by_name('output_alpha:0'),
                                         feed_dict={'input_image:0': roi_img[:, :, ::-1]})
            head_alpha = np.zeros((h, w))
            head_alpha[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha[:, :, 2]
            all_head_alpha.append(head_alpha)

        print('all_head_alpha', all_head_alpha)
        # return all_head_alpha[0]



    def detect_face(self, img):
        h, w, c = img.shape
        input_img = cv2.resize(img[:, :, ::-1], (512, 512))
        boxes, scores, num_detections = self.sess_detect.run(
            [self.sess_detect.graph.get_tensor_by_name('tower_0/boxes:0'),
             self.sess_detect.graph.get_tensor_by_name('tower_0/scores:0'),
             self.sess_detect.graph.get_tensor_by_name('tower_0/num_detections:0')],
            feed_dict={'tower_0/images:0': input_img[np.newaxis], 'training_flag:0': False})
        faceRects = []
        for i in range(num_detections[0]):
            if scores[0, i] < 0.5:
                continue
            y1 = np.int(boxes[0, i, 0] * h)
            x1 = np.int(boxes[0, i, 1] * w)
            y2 = np.int(boxes[0, i, 2] * h)
            x2 = np.int(boxes[0, i, 3] * w)
            if x2 <= x1 + 3 or y2 <= y1 + 3:
                continue
            faceRects.append((y1, y2, x1, x2, y2 - y1, x2 - x1))
        sorted(faceRects, key=lambda x: x[4] * x[5], reverse=True)
        return faceRects

    def pad_box(self, y1, y2, x1, x2, left_ratio, right_ratio, top_ratio, bottom_ratio, h, w):
        box_w = x2 - x1
        box_h = y2 - y1
        pad_y1 = np.maximum(np.int(y1 - top_ratio * box_h), 0)
        pad_y2 = np.minimum(np.int(y2 + bottom_ratio * box_h), h - 1)
        pad_x1 = np.maximum(np.int(x1 - left_ratio * box_w), 0)
        pad_x2 = np.minimum(np.int(x2 + right_ratio * box_w), w - 1)
        return pad_y1, pad_y2, pad_x1, pad_x2



if __name__ == "__main__":
    img = cv2.imread('12345/images/0001.jpg')
    print(img.shape)
    fp = human_segmenter(model_path='assets/matting_human.pb')

    rgba = fp.run(img)
    # cv2.imwrite("human_mask1.png",mask)
    print("test done")