现成的给数据集添加口罩的代码

添加的口罩模板一定要是用ＰＮＧ，ＲＧＢＡ四通道，转换图片到PNG的代码就不放了，找的开源的github代码，根据需要小改了下，可用

class FaceWithMask(object):

    def mask_template(self,):
        current_path = os.path.abspath('../')
        dir = current_path  + '/data/images'
        mask_list = os.listdir(dir)
        mask_list = [os.path.join(dir,x) for x in mask_list]
        mask = np.random.choice(mask_list)
        return mask

    @staticmethod
    def get_distance_from_point_to_line(point, line_point1, line_point2):
        distance = np.abs((line_point2[1] - line_point1[1]) * point[0] +
                          (line_point1[0] - line_point2[0]) * point[1] +
                          (line_point2[0] - line_point1[0]) * line_point1[1] +
                          (line_point1[1] - line_point2[1]) * line_point1[0]) / \
                   np.sqrt((line_point2[1] - line_point1[1]) * (line_point2[1] - line_point1[1]) +
                           (line_point1[0] - line_point2[0]) * (line_point1[0] - line_point2[0]))
        return int(distance)

    def add_mask(self,image_ori_,image_crop_ori,mask_):
        width, height = 128,128
        image_crop_ = cv2.resize(np.array(image_crop_ori), (width, height))
        images = Image.fromarray(image_crop_.astype('uint8')).convert('L')
        images = np.expand_dims(images, axis=0)
        images = np.expand_dims(images, axis=0)
        images = torch.from_numpy(images)
        images = images.float()
        net = face_landmark68()
        res = net(images).cpu().detach().numpy()

        face_landmark = np.array(res[0][0:136]).reshape([-1, 2])
        face_landmark[:, 0] *= width*0.5
        face_landmark[:, 1] *= height*0.5

        # import matplotlib.pyplot as plt
        # plt.imshow(image_crop_)
        # plt.scatter(face_landmark[:,0],face_landmark[:,1],s=15,c='y')
        # plt.show()

        nose_point = face_landmark[28]
        nose_v = np.array(nose_point)

        chin_bottom_point = face_landmark[8]
        chin_bottom_v = np.array(chin_bottom_point)
        chin_left_point = face_landmark[2]
        chin_right_point = face_landmark[14]

        # split mask and resize
        width = mask_.width
        height = mask_.height
        width_ratio = 1.2
        new_height = int(np.linalg.norm(nose_v - chin_bottom_v))

        # left
        mask_left_img = mask_.crop((0, 0, width // 2, height))
        mask_left_width = self.get_distance_from_point_to_line(chin_left_point, nose_point, chin_bottom_point)
        mask_left_width = int(mask_left_width * width_ratio)
        mask_left_img = mask_left_img.resize((mask_left_width, new_height))

        # right
        mask_right_img = mask_.crop((width // 2, 0, width, height))
        mask_right_width = self.get_distance_from_point_to_line(chin_right_point, nose_point, chin_bottom_point)
        mask_right_width = int(mask_right_width * width_ratio)
        mask_right_img = mask_right_img.resize((mask_right_width, new_height))

        # merge mask
        size = (mask_left_img.width + mask_right_img.width, new_height)
        mask_img = Image.new('RGBA', size)
        mask_img.paste(mask_left_img, (0, 0), mask_left_img)
        mask_img.paste(mask_right_img, (mask_left_img.width, 0), mask_right_img)

        # rotate mask
        angle = np.arctan2(chin_bottom_point[1] - nose_point[1], chin_bottom_point[0] - nose_point[0])
        rotated_mask_img = mask_img.rotate(angle, expand=True)

        # calculate mask location
        center_x = (nose_point[0] + chin_bottom_point[0]) // 2
        center_y = (nose_point[1] + chin_bottom_point[1]) // 2

        offset = mask_img.width // 2 - mask_left_img.width
        radian = angle * np.pi / 180
        box_x = center_x + int(offset * np.cos(radian)) - rotated_mask_img.width // 2
        box_y = center_y + int(offset * np.sin(radian)) - rotated_mask_img.height // 2

        # add mask
        # image_crop_ = Image.fromarray(image_crop_ori.astype('uint8'))
        # image_crop_ = Image.fromarray(image_crop_.astype('uint8'))
        image_crop_ori.paste(mask_img, (int(box_x), int(box_y)), mask_img)
        # image_crop_.paste(mask_img, (int(box_x), int(box_y)), mask_img)

        image_crop_ori = np.array(image_crop_ori)

        # import matplotlib.pyplot as plt
        # plt.imshow(image_crop_ori)
        # plt.show()

        return image_crop_ori