前言
随着 Hugging Face 的爆火，平台上出现了越来越多的有趣的模型和数据集，目前仅模型数量就高达 4 万 5 千多个。
这些模型有一个有趣的特点，在云平台上跑的好好的，但是一旦想在本地跑起来就得各种“费劲”折腾，项目关联的 GitHub 中总是能看到用户反馈：这个模型和代码，我本地跑不起来，运行环境和调用代码搞起来太麻烦了。
在日常的工作和学习中，我们也会经常遇到类似上面 Hugging Face 的情况：许多模型在“云上”跑的好好的，但是一到本地就跑不起来了，这或许是因为“操作系统环境、设备 CPU 架构（x86/ ARM）差异”、或许是因为“Python 运行时版本过高或过低”、或许是因为“某个 PIP 安装的软件包版本不对”、“冗长的示例代码中写死了一堆东西”…
用 Docker 容器搭配 Towhee，制作模型的一键运行环境。以腾讯 ARC 实验室开源的 GFPGAN 模型为例。
制作 PyTorch 模型使用的通用 Docker 基础镜像
通过容器解决或避免以下的情况
- 想要避免不同项目之间的环境造成干扰（污染）
- 想要确保项目依赖清晰，任何人都能够在任何设备上复现结果
- 想要复现模型的时间成本更低一些，不喜欢折腾 80% 重复的模型调优之外的工作内容（尤其是环境、基础配置）
考虑到模型可能需要在 x86 和 ARM 两类设备上运行，推荐使用 miniconda3 这个基于 debian 内置了 conda 工具包的基础镜像。
Dockerfile
FROM continuumio/miniconda3:4.11.0
LABEL maintainer=waringid@gmail.com
ARG USE_MIRROR=1
RUN if [ "$USE_MIRROR" = true ] ; then sed -i -e "s/deb.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && sed -i -e "s/security.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list; fi
RUN apt update && \
apt install -y libgl1-mesa-glx && \
apt autoremove -y
RUN if [ "$USE_MIRROR" = true ] ; then pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/; fi
RUN conda install -y pytorch
RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0
RUN pip install towhee
docker build -t waringid/
制作具体模型的应用镜像
下载所需的模型文件
考虑到国内网络下载 Hugging Face 和 GitHub 模型比较慢，还容易出现网络中断。我推荐大家在做应用模型构建的时候，可以考虑提前进行依赖模型的下载，在构建镜像的过程中，将模型放置到合适的目录位置即可。模型下载可以参考 D004-使用 llama.cpp 转换模型程序的模型下载方法，当然也可以单独下载。
在 GFPGAN 项目中一共依赖俩模型文件，一个是 https://github.com/xinntao/facexlib 项目中基于 ResNet50 的人脸检测模型，另一个是用于图片修复的 GFPGAN 对抗网络模型，也就是传统意义上的“主角”。
第一个模型文件 detection_Resnet50_Final.pth ，可以在 https://github.com/xinntao/facexlib/releases/tag/v0.1.0 中获取；第二个模型则需要我们根据自己的设备状况，来做具体选择：
- 如果你需要使用 CPU 来跑模型，可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.2.0 中下载 GFPGANCleanv1-NoCE-C2.pth；或者在https://github.com/TencentARC/GFPGAN/releases/tag/v1.3.0中下载 GFPGANv1.3.pth，这类模型可以完成黑白人像图片的处理。
- 如果你可以使用 GPU 来跑模型，可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.1.0 中下载 GFPGANv1.pth；或者在 https://share.weiyun.com/ShYoCCoc 中进行模型文件下载，这类模型可以处理带颜色的人像图片。
- 除了 GitHub 之外，我们也可以选择直接从 Hugging Face 下载模型（只是可选版本不像上面那么多）：https://huggingface.co/TencentARC/GFPGANv1/tree/main
python 编写模型调用程序
在 GFPGAN 项目中可以找到官方的模型使用示例：https://github.com/TencentARC/GFPGAN/blob/master/inference_gfpgan.py。
app.py
from gfpgan import GFPGANer
import towhee
@towhee.register
class GFPGANerOp:
def __init__(self,
model_path='/GFPGAN.pth',
upscale=2,
arch='clean',
channel_multiplier=2,
bg_upsampler=None) -> None:
self._restorer = GFPGANer(model_path, upscale, arch, channel_multiplier, bg_upsampler)
def __call__(self, img):
cropped_faces, restored_faces, restored_img = self._restorer.enhance(
img, has_aligned=False, only_center_face=False, paste_back=True)
return restored_faces[0][:, :, ::-1]
(
towhee.glob['path']('*.jpg')
.image_load['path', 'img']()
.GFPGANerOp['img','face']()
.show(formatter=dict(img='image', face='image'))
Dockerfile
将下载好的模型文件和新的 Dockerfile 文件放置于相同目录之后，通过整合 Dockerfile 的内容，完成项目依赖的安装，并将模型放置在容器内合适的目录位置。
FROM waringid/docker-pytorch-playground
RUN pip install gfpgan==1.3.8 realesrgan==0.3.0 facexlib==0.3.0 gradio==3.39.0
COPY detection_Resnet50_Final.pth /opt/conda/lib/python3.9/site-packages/facexlib/weights/detection_Resnet50_Final.pth
# 尺寸大一些的模型文件，可以选择使用挂载的方式
# 而不在此处直接 COPY 到容器内部
COPY GFPGANCleanv1-NoCE-C2.pth /GFPGAN.pth
COPY app.py /entrypoint.py
WORKDIR /data
RUN pip install IPython pandas
RUN sed -i -e "s/display(HTML(table))/with open('result.html', 'w') as file:\n file.write(HTML(table).data)/" /opt/conda/lib/python3.9/site-packages/towhee/datacollection/mixins/display.py
CMD ["python3", "/entrypoint.py"]
docker build -t pytorch-playground-gfpgan -f Dockerfile .
模型应用镜像的使用
如果上一步已经下载了模型文件，并将模型文件打包到了镜像中，那么我们只需要下载一些黑白或者彩色的包含人像的图片（根据模型来选择），将它们放在一个目录中（比如 data目录），然后执行一行命令就能够完成模型的调用啦
docker run --rm -it -v `pwd`/model/GFPGANCleanv1-NoCE-C2.pth:/GFPGAN.pth -v `pwd`/data:/data pytorch-playground-gfpgan
如果在上文构建应用模型镜像时，没有选择将 GFPGAN 模型打包到镜像中，那么我们就需要使用文件挂载的方式，来运行模型了。为了项目结构的清晰，可以在项目中创建 model 的目录，来存放上文中提到的模型文件。请忽略 Dockerfile 和 app.py 这几个文件。
在线版的图像模型应用
前面的内容是最简单的模型集成方式，以下的内容通过提供在线访问的方式实现不同的模型和参数的调用，能够在 web 页面上传需要优化的图片并且能够在线展示模型优化后的文件。
先按以下的目录格式下载保存相关的模型（模型文件比较大，提前准备。可以直接下载网盘的资源 https://pan.baidu.com/s/1j6JyPEpCOM4D9NYUnPh0SQ?pwd=soul）
app.py
import os
import cv2
import gradio as gr
import torch
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer
def set_realesrgan(version):
half=True if torch.cuda.is_available() else False
if version == 'v3':
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
return RealESRGANer(scale=4, model_path='model/realesr-general-x4v3.pth', model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
else:
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
return RealESRGANer(scale=2, model_path='model/RealESRGAN_x2plus.pth', model=model, tile=400, tile_pad=10, pre_pad=0, half=half)
def set_face_enhancer(upsampler, version):
if version == 'RestoreFormer':
return GFPGANer(model_path='model/RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
else:
if version == 'CPU':
model_path='model/GFPGANCleanv1-NoCE-C2.pth'
elif version == 'v1':
arch = 'original'
channel_multiplier = 1
model_path='model/GFPGANv1.pth'
elif version == 'v1.2':
arch = 'clean'
channel_multiplier = 2
model_path='model/GFPGANv1.2.pth'
elif version == 'v1.3':
arch = 'clean'
channel_multiplier = 2
model_path='model/GFPGANv1.3.pth'
elif version == 'v1.4':
arch = 'clean'
channel_multiplier = 2
model_path='model/GFPGANv1.4.pth'
return GFPGANer(model_path=model_path, upscale=2, arch=arch, channel_multiplier=channel_multiplier, bg_upsampler=upsampler)
os.makedirs('output', exist_ok=True)
def inference(img, realesrgan_version, gfpgan_version, scale):
if scale > 4:
scale = 4
elif scale < 0:
scale = 1
try:
extension = os.path.splitext(os.path.basename(str(img)))[1]
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
if len(img.shape) == 3 and img.shape[2] == 4:
img_mode = 'RGBA'
elif len(img.shape) == 2: # for gray inputs
img_mode = None
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
else:
img_mode = None
h, w = img.shape[0:2]
if h > 3500 or w > 3500:
print('too large size')
return None, None
if h < 300:
img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
upsampler = set_realesrgan(realesrgan_version)
face_enhancer = set_face_enhancer(upsampler, gfpgan_version)
try:
_, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=None)
except RuntimeError as error:
print('Error', error)
try:
if scale != 2:
interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
h, w = img.shape[0:2]
output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
except Exception as error:
print('wrong scale input.', error)
if img_mode == 'RGBA': # RGBA images should be saved in png format
extension = 'png'
else:
extension = 'jpg'
save_path = f'output/out.{extension}'
cv2.imwrite(save_path, output)
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
return output, save_path
except Exception as error:
print('global exception', error)
return None, None
description = r"""Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>GFPGAN: Towards Real-World Blind Face Restoration with Generative Facial Prior</a>.<br>
这个应用可以被用来修复你的**老照片**或者改善**AI生成的人脸**。<br>
你只需要上传你的图片，就可以使用它。<br>
如果你觉得这个项目帮助到了你，不妨为 <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>它</a> 点个Star吧 :-D <br>
"""
radio_realesrgan_version = gr.Radio(['v3'], type="value", value='v3', label='RealESR GAN version')
radio_gfpgan_version = gr.Radio(['CPU', 'v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], type="value", value='v1.4', label='version')
if os.environ.get('BASICSR_JIT') is not None:
radio_realesrgan_version = gr.Radio(['v2', 'v3'], type="value", value='v3', label='RealESR GAN version')
radio_gfpgan_version = gr.Radio(['CPU', 'v1', 'v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], type="value", value='v1.4', label='version')
app = gr.Interface(
inference, [
gr.Image(type="filepath", label="Input"),
radio_realesrgan_version,
radio_gfpgan_version,
gr.Number(label="Rescaling factor", value=2),
], [
gr.Image(type="numpy", label="Output (The whole image)"),
gr.File(label="Download the output image")
],
title="GFPGAN: 实用的人脸修复算法",
description=description,
article="<p style='text-align: center'>written by: <a href='https://wiki.waringid.me' target='_blank'></a></p>", # examples=[['AI-generate.jpg', 'v1.4', 2], ['lincoln.jpg', 'v1.4', 2], ['Blake_Lively.jpg', 'v1.4', 2], ['10045.png', 'v1.4', 2]]
)
app.queue()
app.launch(server_name="0.0.0.0")
Dockerfile
FROM nvcr.io/nvidia/pytorch:23.04-py3 LABEL org.opencontainers.image.authors="waringid@gmail.com" RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ RUN pip install gfpgan==1.3.8 realesrgan==0.3.0 facexlib==0.3.0 gradio==3.39.0 WORKDIR /app COPY ./app.py ./ CMD ["python", "app.py"]
docker build -t docker-gfpgan . -f docker/Dockerfile
模型应用
docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --rm -it \ -v `pwd`/model:/app/model -v `pwd`/gfpgan:/app/gfpgan -p 7860:7860 docker-gfpgan
添加评论