AI一键去纹身，还原明星「庐山真面目」

Python学习与数据挖掘

共 5160字，需浏览 11分钟

· 2021-04-15

这两天在 Reddit 上看到了一个挺好玩的算法，立马连夜写文过来分享一波。

AI去除纹身：

国内的一些影视作品、综艺节目、B站视频等，都有一些去除纹身的需求。

有些时候，我们需要把一些人身上的纹身覆盖掉，以避免引人效仿。

来自印度的机器学习研究者 Vijish Madhavan 最近开源的一个机器学习工具 SkinDeep 满足了我们这个需求。

项目灵感来源于贾斯汀 · 比伯一次拍 MV ，化妆师花费几个小时的时间才把一身的纹身覆盖掉，采用图像的处理方式可以更省事省力。

SkinDeep

项目地址：https://github.com/vijishmadhavan/SkinDeep

使用 SkinDeep 相比于 Photoshop 更省事一些，毕竟节省了几个小时的时间去修修补补。

我们先看下效果，下图中第一行为输入图像，第二行为输出图像，输出结果明显感觉到，纹身被去除了。

脸部有大量密集纹身的图像，还有其他装饰，AI 的纹身去除效果也非常好：

与专业图像处理软件 photoshop 相比，效果也不错：

看起来 SkinDeep 的效果还不错，但如果纹身是彩色的，还会有一些痕迹。

数据集

完成这个项目需要大量的图像对，因为没有合适的数据集，很多时候训练内容采用合成数据来完成，具体来说：

首先将 APDrawing 数据集图像对与一些背景去掉纹身设计的图像叠加在一起，使用 Python OpenCV 实现；
绘制数据集模拟纹身线条，有助于模型学习和删除这些线条；
APDrawing 数据集只有头像，对于全身图像，项目作者采用了以前的项目 ArtLine，并将输出与输入图像叠加在一起；

ImageDraw.Draw 与森林绿色（forest green colour）色码一起使用，并随机放置在身体图像上，类似于 fast.ai 中的 Crappify ；
Photoshop 也被用来在需要弯曲和角度改变的对象上放置纹身。

简单看了下代码，实现很简单，采用 fast.ai 封装的 unet 进行端到端的训练，并不复杂。

算法测试

开源代码使用 ipynb 文件，测试方法很简单，就两种：

本地：安装 fast.ai ，用 jupyter 打开 ipynb。
colab：直接用 colab 打开。

本地搭建环境，所需要安装的依赖：

fastai==1.0.61
numpy==1.17.2
pandas==1.1.2 
torch==1.6.0
torchvision===0.7.0

本地用 Anaconda 安装好第三方库，下载好代码后，用 jupyter 打开 ipynb 运行。

代码并不复杂，核心代码，就这几行：

import fastai
from fastai.vision import *
from fastai.utils.mem import *
from fastai.vision import open_image, load_learner, image, torch
import numpy as np
import urllib.request
import PIL.Image
from io import BytesIO
import torchvision.transforms as T
from PIL import Image
import requests
from io import BytesIO
import fastai
from fastai.vision import *
from fastai.utils.mem import *
from fastai.vision import open_image, load_learner, image, torch
import numpy as np
import urllib.request
import PIL.Image
from io import BytesIO
import torchvision.transforms as T

class FeatureLoss(nn.Module):
    def __init__(self, m_feat, layer_ids, layer_wgts):
        super().__init__()
        self.m_feat = m_feat
        self.loss_features = [self.m_feat[i] for i in layer_ids]
        self.hooks = hook_outputs(self.loss_features, detach=False)
        self.wgts = layer_wgts
        self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))
              ] + [f'gram_{i}' for i in range(len(layer_ids))]

    def make_features(self, x, clone=False):
        self.m_feat(x)
        return [(o.clone() if clone else o) for o in self.hooks.stored]
    
    def forward(self, input, target):
        out_feat = self.make_features(target, clone=True)
        in_feat = self.make_features(input)
        self.feat_losses = [base_loss(input,target)]
        self.feat_losses += [base_loss(f_in, f_out)*w
                             for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
        self.feat_losses += [base_loss(gram_matrix(f_in), gram_matrix(f_out))*w**2 * 5e3
                             for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
        self.metrics = dict(zip(self.metric_names, self.feat_losses))
        return sum(self.feat_losses)
    
    def __del__(self): self.hooks.remove()