openvino系列 14. GMCNN 图像修复

openvino系列 14. GMCNN 图像修复

此案例演示了如何在 OpenVINO 中使用图像修复模型。 我们使用来自 Open Model Zoo 的 GMCNN 模型。该模型能够在给定篡改图像的情况下进行图像修复。

此案例实现的流程：

1. 下载模型gmcnn-places2-tf以及中间件IR转换；
2. 加载IR模型；
3. 导入测试图片，新建mask，两者合成为mask测试图片；
4. 调整mask以及mask测试图片大小，使之与模型输入大小匹配；
5. 模型推理，显示结果。

环境描述：

• 本案例运行环境：Win10，10代i5笔记本
• IDE：VSCode
• openvino版本：2022.1
• 代码链接，10-image-inpainting

---

---

1 关于预训练模型

1.1 关于GMCNN

GMCNN源于论文Image Inpainting via Generative Multi-column Convolutional Neural Neworks。这里对算法逻辑进行一个大致的解释。

此论文对图像修复过程中三个方面进行了优化：（1）提取图片合适的特征 （2）寻找相似的块 （3）综合辅助信息

1. 提取图片合适的特征：提出了一种多尺度的CNN结构，因为多尺度结构可以将图像分解成具有不同感受野和特征分辨率的分量；
2. 寻找相似的块：提出了一种隐式多样化马尔可夫随机场（ID-MRF）项；
3. 综合辅助信息：设计了一种新的置信驱动的重建损失，根据空间位置约束生成内容。

模型的训练是一种端到端的方式，输入是X破损的图片和 mask M，缺损的区域的填充值为0，M是二进制掩码，0 代表已知的像素，1代表破损区域。模型的测试阶段仅仅只有生成网络（多尺度的CNN结构网络）被使用。

如上图所示，这个算法的网络架构分成三个部分：

1. 生成网络（多尺度的CNN结构）；
2. 全局和局部鉴别器网络；
3. 预训练的VGG网络来计算ID-MRF loss。

这里，我们不对论文的具体内容进行解读，而更专注于如何基于 OpenVINO 使用这个算法。

1.2 GMCNN 预训练模型

OpenVINO “收录了”这个GMCNN的TensorFlow预训练模型。这里我们对这么模型的输入输出进行一个描述：

	GMCNN TensorFlow 预训练模型	GMCNN IR 模型
输入：Mask Image	[1,512,680,3]，对应 [B,H,W,C]	[1,512,680,3]，对应 [B,H,W,C]
输入：Mask	[1,512,680,1]，对应 [B,H,W,C]	[1,512,680,1]，对应 [B,H,W,C]
输出	[1,512,680,3]，对应 [B,H,W,C]	[1,3,512,680]，对应 [B,C,H,W]

B - batch size；H - image height；W - image width；C - number of channels。

也就是说，这个算法的输入有两张图，一张是 Mask，另一张是打了 Mask 的照片。

2 代码实现

2.1 下载模型以及中间件IR转换

我们下载下来的模型是 TensorFlow pb 格式的，所以需要将其转化为 IR 中间件格式。IR（中间表示）模型由一个.xml文件（包含有关网络拓扑的信息）和一个.bin文件（包含权重和偏差二进制数据）组成。 read_model()函数会读取IR模型。我们一般这两个文件（.xml和.bin）放于同一目录中，并且具有相同的文件名。

代码如下：

import sys
from pathlib import Path
import cv2
import matplotlib.pyplot as plt
import numpy as np
from openvino.runtime import Core

# Directory where model will be downloaded
base_model_dir = "model"
# Model name as named in Open Model Zoo
model_name = "gmcnn-places2-tf"

model_path = Path(f"{base_model_dir}/public/{model_name}/{model_name}/frozen_model.pb")
if not model_path.exists():
    download_command = f"omz_downloader " \
                       f"--name {model_name} " \
                       f"--output_dir {base_model_dir}"
    ! $download_command
else:
    print("Already downloaded")
print("1 - Download gmcnn-places2-tf model from Open Model Zoo.")

# 模型转换
precision = "FP16"
ir_path = Path(f"{base_model_dir}/public/{model_name}/{precision}/{model_name}.xml")

# Run Model Optimizer if the IR model file does not exist
if not ir_path.exists():
    print("Exporting TensorFlow model to IR... This may take a few minutes.")
    convert_command = f"omz_converter " \
                      f"--name {model_name} " \
                      f"--download_dir {base_model_dir} " \
                      f"--precisions {precision}"
    ! $convert_command
else:
    print("IR model already exists.")
print("2 - Convert gmcnn-places2-tf Tensorflow model into IR format with {} precision.".format(precision))

2.2 加载模型

现在我们将加载 IR 模型。 然后：

1. 初始化 OpenVINO Runtime
2. 从 *.bin 和 *.xml 文件中读取网络
3. 编译模型
4. 获取输入输出节点

代码如下：

print("3 - Load IR model.")
core = Core()
# Read the model.xml and weights file
model = core.read_model(model=ir_path)
# Load the model on to the CPU
compiled_model = core.compile_model(model=model, device_name="CPU")
# Store the input and output nodes
input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)

print("- input 1 (image with mask) layer: {}".format(compiled_model.input(0).shape))
print("- input 2 (mask) layer: {}".format(compiled_model.input(1).shape))
print("- output layer: {}".format(compiled_model.output(0).shape))

2.3 创建 Mask 以及 Mask Image

首先，我们创建一个单通道mask。

print("4 - Create mask, load image and generate mask image.")

def create_mask(image_width, image_height, size_x=30, size_y=30, number=1):
    """
    Create a square mask of defined size on a random location

    :param: image_width: width of the image
    :param: image_height: height of the image
    :param: size: size in pixels of one side
    :returns:
            mask: grayscale float32 mask of size shaped [image_height, image_width, 1]
    """

    mask = np.zeros((image_height, image_width, 1), dtype=np.float32)
    for _ in range(number):
        start_x = np.random.randint(image_width - size_x)
        start_y = np.random.randint(image_height - size_y)
        cv2.rectangle(img=mask,
                      pt1=(start_x, start_y),
                      pt2=(start_x + size_x, start_y + size_y),
                      color=(1, 1, 1),
                      thickness=cv2.FILLED)
    return mask

N, H, W, C = input_layer.shape
print("模型输入尺寸：{}".format(input_layer.shape))
# Generate a square mask of size WxH with number of "holes"
mask = create_mask(image_width=W, image_height=H, size_x=50, size_y=50, number=15)
# This mask will be laid over the input image as noise
plt.figure(figsize=(16, 12))
plt.imshow(cv2.cvtColor(mask, cv2.COLOR_BGR2RGB));

接下来，我们导入图片，并给图片“打码”：

'''
加载和调整图像大小
'''
# Read the image
image = cv2.imread("data/laptop.png")
# Resize image to meet network expected input sizes
resized_image = cv2.resize(src=image, dsize=(W, H), interpolation=cv2.INTER_AREA)
#plt.figure(figsize=(16, 12))
#plt.imshow(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB));
'''
生成mask图像
'''
# Generating masked image
masked_image = (resized_image * (1 - mask) + 255 * mask).astype(np.uint8)
plt.figure(figsize=(16, 12))
plt.imshow(cv2.cvtColor(masked_image, cv2.COLOR_BGR2RGB));

2.4 预处理

调整mask以及mask测试图片大小，使之与模型输入大小匹配；

• masked_image.shape = (512,680,3) -----> 模型期望 = (1,512,680,3)
• resized_mask.shape = (512,680,1) -----> 模型期望 = (1,512,680,1)

print("5 - resize mask image and mask to be the same shape as model inputs")

print("- masked image shape (before reshape): {}".format(masked_image.shape))
print("- mask shape (after reshape): {}".format(mask.shape))

masked_image = masked_image[None, ...]
mask = mask[None, ...]

print("- masked image shape (before reshape): {}".format(masked_image.shape))
print("- mask shape (after reshape): {}".format(mask.shape))

2.5 推理

注意，这里的输入有两个（虽然我不能明白为什么）：mask以及masked_image（官方链接）。 然后显示恢复的图像。

print("6 - Model compile and demo results.")

result = compiled_model([masked_image, mask])[output_layer]
result = result.squeeze().astype(np.uint8)
plt.figure(figsize=(16, 12))
plt.imshow(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))

我们将Terminal中的返回放到一起，就能很清楚地看到代码的整个逻辑：

1 - Download gmcnn-places2-tf model from Open Model Zoo.
IR model already exists.
2 - Convert gmcnn-places2-tf Tensorflow model into IR format with FP16 precision.
3 - Load IR model.
- input 1 (image with mask) layer: {1, 512, 680, 3}
- input 2 (mask) layer: {1, 512, 680, 1}
- output layer: {1, 512, 680, 3}
4 - Create mask, load image and generate mask image.
5 - resize mask image and mask to be the same shape as model inputs
- masked image shape (before reshape): (512, 680, 3)
- mask shape (after reshape): (512, 680, 1)
- masked image shape (before reshape): (1, 512, 680, 3)
- mask shape (after reshape): (1, 512, 680, 1)
6 - Model compile and demo results.