mmgeneration cyclegan translation_demo.py 批量转图片

使用cycleGAN批量将图片转化为梵高风格

最新推荐文章于 2025-12-03 12:28:21 发布

原创最新推荐文章于 2025-12-03 12:28:21 发布 · 298 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#python #计算机视觉 #opencv

该代码示例展示了如何利用预训练的cycleGAN模型将一组图片转换为梵高艺术风格，并保存转换后的图像。通过指定配置文件和预训练权重，模型读取输入图片并生成梵高风格的图像，然后可以遍历目录批量处理所有.jpg和.png文件。

## cycleGAN 常规图片批量转换为梵高domain(或者其他domain)并保存于该文件夹中.


# 导入mmcv和mmgeneration
# import mmcv
from mmgen.apis import init_model, sample_img2img_model
import os
# 导入 opencv
import cv2

# 导入numpy和matplotlib
import numpy as np
import matplotlib.pyplot as plt
# 照片转梵高

# 指定config文件路径
config_file = 'configs/cyclegan/cyclegan_lsgan_resnet_in_facades_b1x1_80k_vangogh2photo.py'

# 指定预训练模型权重文件路径
checkpoint_file = 'work_dirs/experiments/cyclegan_vangogh2photo/ckpt/cyclegan_vangogh2photo/iter_best.pth'
model = init_model(config_file, checkpoint_file)


# 输出图像宽高像素尺寸
img_size = 1024


def gen_vis_fake_img(input_path, model, target_domain='vangogh', figsize=60, save_path=None):
    # 读入输入图像，获取高宽尺寸  folder = 'Picture/'
    # ch_filepath =  '1_'+filepath
    input_img = cv2.imread(input_path)
    # 生成梵高油画图像，注意 target_domain 要设置正确
    fake_imgs = sample_img2img_model(model, input_path, target_domain=target_domain)
    # 获取生成图像尺寸
    img_size = fake_imgs.shape[2]

    # 分别抽取RGB三通道图像，归一化为0-255的uint8自然图像
    RGB = np.zeros((img_size, img_size, 3))
    RGB[:, :, 0] = fake_imgs[0][2]
    RGB[:, :, 1] = fake_imgs[0][1]
    RGB[:, :, 2] = fake_imgs[0][0]
    # 将生成图转为输入图像大小
    RGB = cv2.resize(RGB, dsize=(input_img.shape[1], input_img.shape[0]))
    # 像素值归一化
    RGB = 255 * (RGB - RGB.min()) / (RGB.max() - RGB.min())
    # 像素值转为整数
    RGB = RGB.astype('uint8')

    if save_path:
        # 导出生成的图像文件
        cv2.imwrite(save_path, cv2.cvtColor(RGB, cv2.COLOR_BGR2RGB))

#以上代码 部分借鉴同济子豪兄.



def traverse_dir(path):
    for file in os.listdir(path):
        file_path = os.path.join(path,file)
        if os.path.isfile(file_path):
            if os.path.splitext(file_path)[1] in ['.jpg','.png']:
                out_file_name = os.path.splitext(file_path)\[0]+'_gan'+os.path.splitext(file_path)[1]
                gen_vis_fake_img(file_path, model, target_domain='vangogh', figsize=60, save_path=out_file_name)
        elif os.path.isdir(file_path):
            traverse_dir(file_path)

if __name__ == '__main__':
    path = './data_pic'
    traverse_dir(path)