Segment-anything更改证件照底色

Meta 开源万物可分割 AI 模型：segment anything model (SAM)。
本文利用SAM进行证件照底色更换（目前不支持尺寸选择）
paper：arxiv.org/abs/2304.02643
code：https://github.com/facebookresearch/segment-anything

一、思路介绍

加载图片image→获取人物mask→将mask区域的图片裁剪出来→更改背景色

二、代码详解（deploy.py）

import numpy as np
import matplotlib.pyplot as plt
import cv2
import syssys.path.append("..")
from segment_anything import sam_model_registry, SamPredictorclass Mask():def getMask(image,point=[300,375]):r''':param image:输入的图像:param point:进行识别的 point:return : 人物 mask'''sam_checkpoint = "checkpoint/sam_vit_l_0b3195.pth"model_type = "vit_l"device = "cuda"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)sam.to(device=device)predictor = SamPredictor(sam)predictor.set_image(image)input_point = np.array([point])input_label = np.array([1])masks, _, _ = predictor.predict(point_coords=input_point,point_labels=input_label,multimask_output=True,)return masks[2]def saveMask(mask,path):r''':param mask:人物 mask:param path:mask 的保存路径'''# 将 numpy.ndarray 数组转换为 OpenCV 图像mask = cv2.cvtColor(np.uint8(mask * 255), cv2.COLOR_RGB2BGR)cv2.imwrite(path, mask)class Image():def getImage(path):r''':param path: image 的路径:return: image'''image = cv2.imread(path)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)return imagedef cutImage(mask_path,image):r''':param mask_path: mask 图片路径:param image: 原始图片:return: 裁剪之后的图像'''# 加载不规则掩码mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)# 将掩码转换为二值图像_, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)# 使用掩码裁剪图像cropped_image = cv2.bitwise_and(image, image, mask=mask)return cropped_imagedef changeBg(bgColor,image,cropped_image,mask):r''':param bgColor: 要修改的背景颜色:param image: 原始图片:param cropped_image: 裁剪之后的图片:param mask: mask 图片:return: 换底色的图片'''# 创建红色背景图像red_bgr = np.zeros_like(image)red_bgr[:] = bgColor# 将裁剪出来的图像放在红色背景上result = red_bgr.copy()result[mask != 0] = cropped_image[mask != 0]return resultif __name__ == '__main__':# 获取原始图片路径image_path = 'notebooks/images/wy.jpg'# 保存的 mask 路径，需要保存一下进行格式转换mask_path = 'notebooks/images/mask.jpg'# RGB（要更改的背景色）color_bg = [223,123,132]assert color_bg is not None,"The background color is empty"# 获取图片image = Image.getImage(path=image_path)h, w, _ = image.shape# → x轴 ;↓ y轴 ;point =>(x,y)point = [int(w/2),int(h/2)]# 加载maskmask = Mask.getMask(image,point)Mask.saveMask(mask,mask_path)# 将mask区域的图片裁剪出来cropped_image = Image.cutImage(mask_path,image)# 更改背景色result = Image.changeBg(color_bg,image= image,cropped_image=cropped_image,mask=mask)plt.figure(figsize=(30, 30))plt.imshow(result)plt.axis('off')plt.show()

项目将上传到csdn。

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！