解锁CANN仓库核心能力：50行代码搭建国产化AIGC图片风格迁移神器

# 安装CANN仓库核心依赖（从官方仓库拉取）
pip install ascend-cann-acl==7.0.0 opencv-python==4.8.1.78 numpy==1.24.4

# 下载预训练的FastStyleTransfer模型（已转换为OM格式）
wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/samples/fast_style_transfer_vangogh.om

import acl
import cv2
import numpy as np

# -------------------------- 核心配置 --------------------------
DEVICE_ID = 0
OM_MODEL_PATH = "./fast_style_transfer_vangogh.om"  # 梵高风格迁移OM模型
INPUT_IMG_PATH = "./test_photo.jpg"  # 输入普通照片
OUTPUT_IMG_PATH = "./styled_photo.jpg"  # 输出风格化照片
IMG_SIZE = (256, 256)  # 模型输入尺寸

# -------------------------- 初始化CANN环境 --------------------------
def init_cann():
    acl.init()
    acl.rt.set_device(DEVICE_ID)
    context, _ = acl.rt.create_context(DEVICE_ID)
    stream, _ = acl.rt.create_stream()
    # 加载OM模型
    model_id, _ = acl.mdl.load_from_file(OM_MODEL_PATH)
    model_desc = acl.mdl.create_desc()
    acl.mdl.get_desc(model_desc, model_id)
    return context, stream, model_id, model_desc

# -------------------------- 图片风格迁移核心逻辑 --------------------------
def style_transfer(img_path, model_id, model_desc, stream):
    # 1. 图片预处理（CANN兼容算子加速）
    img = cv2.imread(img_path)
    img_resized = cv2.resize(img, IMG_SIZE)  # 缩放
    img_normalized = (img_resized / 255.0).astype(np.float32)  # 归一化
    img_transposed = np.transpose(img_normalized, (2, 0, 1))  # HWC→CHW
    input_data = np.expand_dims(img_transposed, axis=0)  # 增加batch维度

    # 2. 准备模型输入/输出内存（CANN核心内存管理）
    input_ptr = acl.util.numpy_to_ptr(input_data)
    input_dataset = acl.mdl.create_dataset()
    input_buffer = acl.create_data_buffer(input_ptr, input_data.nbytes)
    acl.mdl.add_dataset_buffer(input_dataset, input_buffer)

    output_dataset = acl.mdl.create_dataset()
    output_size = acl.mdl.get_output_size_by_index(model_desc, 0)
    output_ptr = acl.rt.malloc(output_size, acl.rt.malloc_mem_type.MEM_DEVICE)
    output_buffer = acl.create_data_buffer(output_ptr, output_size)
    acl.mdl.add_dataset_buffer(output_dataset, output_buffer)

    # 3. 调用CANN推理（核心步骤）
    acl.mdl.execute_async(model_id, input_dataset, output_dataset, stream)
    acl.rt.synchronize_stream(stream)

    # 4. 后处理：转换为可视化图片
    output_data = np.empty((1, 3, IMG_SIZE[0], IMG_SIZE[1]), dtype=np.float32)
    acl.rt.memcpy(
        acl.util.numpy_to_ptr(output_data),
        output_data.nbytes,
        output_ptr,
        output_size,
        acl.rt.memcpy_kind.MEMCPY_DEVICE_TO_HOST
    )
    # 反归一化+维度转换
    styled_img = np.squeeze(output_data)  # 去除batch维度
    styled_img = np.transpose(styled_img, (1, 2, 0))  # CHW→HWC
    styled_img = (styled_img * 255).clip(0, 255).astype(np.uint8)  # 反归一化

    # 5. 释放临时资源
    acl.destroy_data_buffer(input_buffer)
    acl.destroy_data_buffer(output_buffer)
    acl.mdl.destroy_dataset(input_dataset)
    acl.mdl.destroy_dataset(output_dataset)
    return styled_img

# -------------------------- 主函数 --------------------------
if __name__ == "__main__":
    try:
        # 初始化+CANN推理
        context, stream, model_id, model_desc = init_cann()
        styled_img = style_transfer(INPUT_IMG_PATH, model_id, model_desc, stream)
        # 保存结果
        cv2.imwrite(OUTPUT_IMG_PATH, styled_img)
        print(f"风格迁移完成！结果已保存至：{OUTPUT_IMG_PATH}")
    except Exception as e:
        print(f"执行失败：{e}")
    finally:
        # 释放所有资源
        acl.mdl.unload(model_id)
        acl.mdl.destroy_desc(model_desc)
        acl.rt.destroy_stream(stream)
        acl.rt.destroy_context(context)
        acl.rt.reset_device(DEVICE_ID)
        acl.finalize()