MMagic实战

从网络下载图像，做Canny 边缘检测后保存

import cv2
import numpy as np
import requests
import io

# 下载图片
url = 'https://th.bing.com/th/id/OIP.2MaWvr94eGVEN1yLHBltqgHaE7?pid=ImgDet&rs=1'
response = requests.get(url)
image = np.asarray(bytearray(response.content), dtype="uint8")
image = cv2.imdecode(image, cv2.IMREAD_GRAYSCALE)

# 保存原始图片
cv2.imwrite('image.jpg', image)

# 使用Canny边缘检测
edges = cv2.Canny(image, threshold1=30, threshold2=100)

# 保存图像
cv2.imwrite('edges.jpg', edges)

处理前后对比

import matplotlib.pyplot as plt

# 读取原图
original = cv2.imread('image.jpg', cv2.IMREAD_COLOR)

# 读取处理后的图像
edges = cv2.imread('edges.jpg', cv2.IMREAD_COLOR)

# OpenCV使用BGR模式，而matplotlib使用RGB模式，所以我们需要转换颜色模式
original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
edges = cv2.cvtColor(edges, cv2.COLOR_BGR2RGB)

# 创建子图来展示对比
fig, axs = plt.subplots(1, 2, figsize=(10, 5))

# 展示原图
axs[0].imshow(original, cmap='gray')
axs[0].set_title('Original Image')

# 展示处理后的图像
axs[1].imshow(edges, cmap='gray')
axs[1].set_title('Edge Image')

# 显示图像
plt.show()

环境配置

!pip3 install openmim
!mim install 'mmcv>=2.0.0'
!mim install 'mmengine'

!git clone https://github.com/open-mmlab/mmagic.git
%cd mmagic
!pip3 install -e .

导入工具包

import cv2
import numpy as np
import mmcv
from mmengine import Config
from PIL import Image

from mmagic.registry import MODELS
from mmagic.utils import register_all_modules

register_all_modules()

载入ControlNet模型

cfg = Config.fromfile('configs/controlnet/controlnet-canny.py')
controlnet = MODELS.build(cfg.model).cuda()

输入Canny边缘图

# control_url = 'https://user-images.githubusercontent.com/28132635/230288866-99603172-04cb-47b3-8adb-d1aa532d1d2c.jpg'
control_img = mmcv.imread('/content/edges.jpg')
control = cv2.Canny(control_img, 100, 200)
control = control[:, :, None]
control = np.concatenate([control] * 3, axis=2)
control = Image.fromarray(control)

咒语Prompt

prompt = 'Please decorate my cottage in a pale blue marine style, with walls painted with lifelike sea creatures!'

执行预测

output_dict = controlnet.infer(prompt, control=control)
samples = output_dict['samples']
for idx, sample in enumerate(samples):
    sample.save(f'sample_{idx}.png')
controls = output_dict['controls']
for idx, control in enumerate(controls):
    control.save(f'control_{idx}.png')

结果展示

import matplotlib.pyplot as plt
import cv2

# 以彩色模式读取图片
original = cv2.imread('/content/mmagic/control_0.png', cv2.IMREAD_COLOR)
edges = cv2.imread('/content/mmagic/sample_0.png', cv2.IMREAD_COLOR)

# OpenCV使用BGR模式，而matplotlib使用RGB模式，所以我们需要转换颜色模式
original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
edges = cv2.cvtColor(edges, cv2.COLOR_BGR2RGB)

# 创建子图来展示对比
fig, axs = plt.subplots(1, 2, figsize=(10, 5))

# 展示原图
axs[0].imshow(original)
axs[0].set_title('Original Image')

# 展示处理后的图像
axs[1].imshow(edges)
axs[1].set_title('Edge Image')

# 显示图像
plt.show()