通过文心大模型生成你的兔年专属美甲

我正在参加「兔了个兔」创意投稿大赛，详情请看：「兔了个兔」创意投稿大赛

  本文章最初的灵感来源于旷视放出来的虚拟试妆，实际上在美妆行业，虚拟试妆也越来越普遍了，恰好手里有指甲分割的数据集，就尝试着玩了玩。使用文心大模型生成自己想要的兔年专属美甲图案，主要基于PaddleSeg套件、文心大模型进行开发，最终的效果如下图所示：

1、技术路线

  考虑到，故在此选择了pp_liteseg作为网络模型解决方案，对指甲部分进行分割，考虑到线上试妆颜色的真实性，不同于其他的图像分割，这里不能将指甲的分割区域进行简单的颜色填充，需要考虑分割部分指甲部位的亮度等真实指甲的特征，本项目中对分割部分的HSV进行调整，准确来说是对饱和度（S）、亮度（V）进行调整，保留真实指甲特征的同时更改颜色，达到虚拟试妆的效果。

1.1 数据集预处理

  根据PaddleSeg中的数据标注方式，在这里，我们需要将数据集进行预处理，将标注区域的像素更改为1，因为此图像分割任务中只包含指甲一个类别。

import numpy as np
import os
import cv2

np.set_printoptions(threshold=np.inf)

images_path = 'data/data166460/Nails/JPEGImages'
masks_path = 'data/data166460/Nails/SegmentationClass'
masked_path = 'data/data166460/Nails/SegmentationClass_masked'

os.makedirs(masked_path)


for img_item in os.listdir(images_path):
    # print(img_item)
    img_path = os.path.join(images_path, img_item)
    img = cv2.imread(img_path)
    mask_path = os.path.join(masks_path, img_item[:-4]+'.png')  # mask是.png格式的
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)  # 将彩色mask以二值图像形式读取
    mask[mask > 0] = 1
    # print(mask)

    cv2.imwrite(os.path.join(masked_path, img_item), mask) 

1.2 数据集探索分析

  探索性数据分析（Exploratory Data Analysis，简称EDA），是指对已有的数据（原始数据）进行分析探索，通过作图、制表、方程拟合、计算特征量等手段探索数据的结构和规律的一种数据分析方法。一般来说，我们最初接触到数据的时候往往是毫无头绪的，不知道如何下手，这时候探索性数据分析就非常有效。

%matplotlib inline

import os
import cv2
import random

from tqdm import tqdm

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt



TRAIN = 'data/data166460/Nails/JPEGImages'

train = os.listdir(TRAIN)

f, axs = plt.subplots(4,5, figsize=(16,8))
for i in range(4):
    for j in range(5):
        axs[i][j].imshow(cv2.cvtColor(cv2.imread(os.path.join(TRAIN, train[random.randint(0,len(train)-1)]), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB));

2、模型训练

  本模型后续部署在小型设备上，故在此选择了一个轻量级网络pp_liteseg_optic_disc_512x512_1k作为分割模型，其配置文件在路径

PaddleSeg/configs/quick_start/pp_liteseg_optic_disc_512x512_1k.yml

  模型整体还是采用跳层连接结构，引入了统一注意力融合模块（空间和通道注意力机制）、金字塔模块来聚合上下文。

2.1 不使用文心大模型的后处理效果

  通过对结果进行后处理，生成AI美妆结果。在这里给出模型预测结果以及经过后处理之后的结果。其中左侧的为经过后处理之后的图片，右侧的为模型输出的结果。通过对比我们可以发现，经过后处理的图片，指甲部分的亮度都被保存，整体的效果要强于模型直接的生成结果。

mask_res_path = 'PaddleSeg/output/result/pseudo_color_prediction/1eecab90-1a92-43a7-b952-0204384e1fae.png'

mask_test_path = 'data/data166460/Nails/SegmentationClass_masked/1eecab90-1a92-43a7-b952-0204384e1fae.png'

img = cv2.imread(ori_path)
mask_img = cv2.imread(mask_test_path, cv2.IMREAD_GRAYSCALE)

img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
img[mask_img > 0,0] = 200
img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)

cv2.imwrite('PaddleSeg/output/result/res.jpg', img)

3.2 使用文心大模型的后处理效果

! pip install wenxin_api
import wenxin_api # 可以通过"pip install wenxin-api"命令安装
from wenxin_api.tasks.text_to_image import TextToImage


wenxin_api.ak = ""
wenxin_api.sk = ""


input_dict = {
    "text": '动漫，兔子，新年' + "超高清，超细节，唯美，插画，壁纸",
    "style": "古风", #解锁更多风格后，非必选参数
    "resolution":"1024*1536" , #也可设置为 1024*1536、1536*1024
    "num": "1",    #功能解锁后，可设置的范围为[1,2,3,4,5,6]
}
rst = TextToImage.create(**input_dict)

!pip install wget
import wget

wget.download(rst['imgUrls'][0], 'source1.jpg')

import cv2
import numpy as np

# 读取图片
img_dir = 'data/data166460/Nails/JPEGImages/1eecab90-1a92-43a7-b952-0204384e1fae.png'
mask_dir = 'data/data166460/Nails/SegmentationClass/1eecab90-1a92-43a7-b952-0204384e1fae.png'

source_dir = 'source.jpg'

img = cv2.imread(img_dir)
mask = cv2.imread(mask_dir)[:,:,0]

# cv2.imwrite('now_img.jpg',img)
# cv2.imwrite('now_mask.jpg',mask*255)

mask = mask.astype('float32')
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype('float32')

source = cv2.imread(source_dir)
source = cv2.resize(source, (hsv.shape[1], hsv.shape[0]))
hsv_source = cv2.cvtColor(source, cv2.COLOR_BGR2HSV).astype('float32')

h_mid = np.median(hsv[mask>0,0])
s_mid = np.median(hsv[mask>0,1])
v_mid = np.median(hsv[mask>0,2])
hsv_mid = np.array([h_mid, s_mid, v_mid]).astype('float32')

hsv[mask>0,0] = (hsv[mask>0,0] - hsv_mid[0] + hsv_source[mask>0,0])
hsv[mask>0,1] = (hsv[mask>0,1] - hsv_mid[1] + hsv_source[mask>0,1]) # *0.5 + hsv[mask>0,1]*0.5

hsv[hsv[:,:,0]>179,0] = 179
hsv[hsv[:,:,0]<0,0] = 0
hsv[hsv[:,:,1]>255,1] = 255
hsv[hsv[:,:,1]<0,1] = 0

result = cv2.cvtColor(hsv.astype('uint8'), cv2.COLOR_HSV2BGR)

cv2.imwrite('result.jpg',result)

3、总结

  项目通过文心大模型、PaddleSeg套件完成了AI美甲，但是还有很多提升的地方，例如指甲边缘的处理，仔细看的话，目前的效果对于指甲边缘的处理不是特别好；其次，是文心大模型生成的图片，可以有更适合美甲的图片，这个需要大家去调整Promot，最后祝大家新春快乐。