阿里开源黑白图片上色算法DDColor的部署与测试并将模型转onnx后用c++推理

阿里开源黑白图片上色算法DDColor的部署与测试并将模型转onnx后用c++推理

文章目录

  • 阿里开源黑白图片上色算法DDColor的部署与测试并将模型转onnx后用c++推理
    • 简介
    • 环境部署
      • 下载源码
      • 安装环境
      • 下载模型
    • 测试一下
    • 看看效果
    • 模型转onnx
    • 测试一下生成的onnx模型
    • 看看效果
    • C++ 推理

简介

DDColor是一种基于深度学习的图像上色技术，它利用卷积神经网络（CNN）对黑白图像进行上色处理。该模型通常包含一个编码器和一个解码器，编码器提取图像的特征，解码器则根据这些特征生成颜色。DDColor模型能够处理多种类型的图像，并生成自然且逼真的颜色效果。它在图像编辑、电影后期制作以及历史照片修复等领域有广泛的应用。

环境部署

下载源码

git clone https://github.com/piddnad/DDColor.git

安装环境

conda create -n ddcolor python=3.9
conda activate ddcolor
pip install -r requirements.txt
python3 setup.py develop
pip install modelscope
pip install onnx
pip install onnxruntime

下载模型

这里下载
或者运行下面的脚本下载：

from modelscope.hub.snapshot_download import snapshot_download
model_dir = snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')
print('model assets saved to %s'%model_dir)
#模型会被下载到modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt

测试一下

import argparse
import cv2
import numpy as np
import os
from tqdm import tqdm
import torch
from basicsr.archs.ddcolor_arch import DDColor
import torch.nn.functional as F

class ImageColorizationPipeline(object):

    def __init__(self, model_path, input_size=256, model_size='large'):

        self.input_size = input_size
        if torch.cuda.is_available():
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')

        if model_size == 'tiny':
            self.encoder_name = 'convnext-t'
        else:
            self.encoder_name = 'convnext-l'

        self.decoder_type = "MultiScaleColorDecoder"

        if self.decoder_type == 'MultiScaleColorDecoder':
            self.model = DDColor(
                encoder_name=self.encoder_name,
                decoder_name='MultiScaleColorDecoder',
                input_size=[self.input_size, self.input_size],
                num_output_channels=2,
                last_norm='Spectral',
                do_normalize=False,
                num_queries=100,
                num_scales=3,
                dec_layers=9,
            ).to(self.device)
        else:
            self.model = DDColor(
                encoder_name=self.encoder_name,
                decoder_name='SingleColorDecoder',
                input_size=[self.input_size, self.input_size],
                num_output_channels=2,
                last_norm='Spectral',
                do_normalize=False,
                num_queries=256,
            ).to(self.device)

        self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))['params'],strict=False)
        self.model.eval()

    @torch.no_grad()
    def process(self, img):
        self.height, self.width = img.shape[:2]
        # print(self.width, self.height)
        # if self.width * self.height < 100000:
        #     self.input_size = 256

        img = (img / 255.0).astype(np.float32)
        orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]  # (h, w, 1)

        # resize rgb image -> lab -> get grey -> rgb
        img = cv2.resize(img, (self.input_size, self.input_size))
        img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
        img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
        img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)

        tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device)

        # (1, 2, self.height, self.width)
        output_ab = self.model(tensor_gray_rgb).cpu()

        # resize ab -> concat original l -> rgb
        output_ab_resize = F.interpolate(output_ab, size=(self.height, self.width))[0].float().numpy().transpose(1, 2, 0)
        output_lab = np.concatenate((orig_l, output_ab_resize), axis=-1)
        output_bgr = cv2.cvtColor(output_lab, cv2.COLOR_LAB2BGR)

        output_img = (output_bgr * 255.0).round().astype(np.uint8)

        return output_img


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str,default='pretrain/net_g_200000.pth')
    parser.add_argument('--input_size', type=int,default=512, help='input size for model')
    parser.add_argument('--model_size', type=str,default='large', help='ddcolor model size')
    args = parser.parse_args()

    colorizer = ImageColorizationPipeline(model_path=args.model_path, input_size=args.input_size, model_size=args.model_size)

    img = cv2.imread("./down.jpg")
    image_out = colorizer.process(img)
    cv2.imwrite("./downout.jpg", image_out)


if __name__ == '__main__':
    main()


python test.py  --model_path=./modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt

看看效果

在这里插入图片描述

在这里插入图片描述
效果看起来非常的nice！

模型转onnx

import argparse
import cv2
import numpy as np
import os
from tqdm import tqdm
import torch
from basicsr.archs.ddcolor_arch import DDColor
import torch.nn.functional as F

class ImageColorizationPipeline(object):

    def __init__(self, model_path, input_size=256, model_size='large'):
        
        self.input_size = input_size
        if torch.cuda.is_available():
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')

        if model_size == 'tiny':
            self.encoder_name = 'convnext-t'
        else:
            self.encoder_name = 'convnext-l'

        self.decoder_type = "MultiScaleColorDecoder"

        if self.decoder_type == 'MultiScaleColorDecoder':
            self.model = DDColor(
                encoder_name=self.encoder_name,
                decoder_name='MultiScaleColorDecoder',
                input_size=[self.input_size, self.input_size],
                num_output_channels=2,
                last_norm='Spectral',
                do_normalize=False,
                num_queries=100,
                num_scales=3,
                dec_layers=9,
            ).to(self.device)
        else:
            self.model = DDColor(
                encoder_name=self.encoder_name,
                decoder_name='SingleColorDecoder',
                input_size=[self.input_size, self.input_size],
                num_output_channels=2,
                last_norm='Spectral',
                do_normalize=False,
                num_queries=256,
            ).to(self.device)

        print(model_path)

        self.model.load_state_dict(
            torch.load(model_path, map_location=torch.device('cpu'))['params'],
            strict=False)
        self.model.eval()

    @torch.no_grad()
    def process(self, img):
        self.height, self.width = img.shape[:2]
        # print(self.width, self.height)
        # if self.width * self.height < 100000:
        #     self.input_size = 256

        img = (img / 255.0).astype(np.float32)
        orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]  # (h, w, 1)

        # resize rgb image -> lab -> get grey -> rgb
        img = cv2.resize(img, (self.input_size, self.input_size))
        img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
        img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
        img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)

        tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device)
        output_ab = self.model(tensor_gray_rgb).cpu()  # (1, 2, self.height, self.width)
        
        # resize ab -> concat original l -> rgb
        output_ab_resize = F.interpolate(output_ab, size=(self.height, self.width))[0].float().numpy().transpose(1, 2, 0)
        output_lab = np.concatenate((orig_l, output_ab_resize), axis=-1)
        output_bgr = cv2.cvtColor(output_lab, cv2.COLOR_LAB2BGR)

        output_img = (output_bgr * 255.0).round().astype(np.uint8)    

        return output_img


    @torch.no_grad()
    def expirt_onnx(self, img):
        self.height, self.width = img.shape[:2]
        
        img = (img / 255.0).astype(np.float32)
        orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]  # (h, w, 1)

        # resize rgb image -> lab -> get grey -> rgb
        img = cv2.resize(img, (self.input_size, self.input_size))
        img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
        img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
        img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)

        tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device)
        
        mymodel = self.model.to('cpu')
        tensor_gray_rgb = tensor_gray_rgb.to('cpu')
        onnx_save_path = "color.onnx"

        torch.onnx.export(mymodel,  # 要导出的模型
                          tensor_gray_rgb,  # 模型的输入
                          onnx_save_path,  # 导出的文件路径
                          export_params=True,  # 是否将训练参数导出
                          opset_version=12,  # 导出的ONNX的操作集版本
                          do_constant_folding=True,  # 是否执行常量折叠优化
                          input_names=['input'],  # 输入张量的名称
                          output_names=['output'],  # 输出张量的名称
                          dynamic_axes={'input': {0: 'batch_size'}, 
                                        'output': {0: 'batch_size'}})
        return



def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, default='pretrain/net_g_200000.pth')
    parser.add_argument('--input_size', type=int, default=512, help='input size for model')
    parser.add_argument('--model_size', type=str, default='large', help='ddcolor model size')
    args = parser.parse_args()

    colorizer = ImageColorizationPipeline(model_path=args.model_path, input_size=args.input_size, model_size=args.model_size)

    img = cv2.imread("./down.jpg")
    image_out = colorizer.expirt_onnx(img)
    # image_out = colorizer.process(img)
    # cv2.imwrite("./downout.jpg", image_out)


if __name__ == '__main__':
    main()


python model2onnx.py  --model_path=./modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt

测试一下生成的onnx模型

import onnxruntime
import cv2
import numpy as np

def colorize_image(input_image_path, output_image_path, model_path):
    input_image = cv2.imread(input_image_path)

    img = (input_image / 255.0).astype(np.float32)
    orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]  # (h, w, 1)
    img = cv2.resize(img, (512, 512))
    img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
    img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
    input_blob = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)

    # Change data layout from HWC to CHW
    input_blob = np.transpose(input_blob, (2, 0, 1))
    input_blob = np.expand_dims(input_blob, axis=0)  # Add batch dimension

    # Initialize ONNX Runtime Inference Session
    session = onnxruntime.InferenceSession(model_path)

    # Perform inference
    output_blob = session.run(None, {'input': input_blob})[0]

    # Post-process the output
    output_blob = np.squeeze(output_blob)  # Remove batch dimension

    # Separate ab channels
    # Change data layout from CHW to HWC
    output_ab = output_blob.transpose((1, 2, 0))

    # Resize to match input image size
    output_ab = cv2.resize(output_ab, (input_image.shape[1], input_image.shape[0]))
    output_lab = np.concatenate((orig_l, output_ab), axis=-1)

    # Convert LAB to BGR
    output_bgr = cv2.cvtColor(output_lab, cv2.COLOR_LAB2BGR)

    output_bgr = output_bgr*255

    # Save the colorized image
    cv2.imwrite(output_image_path, output_bgr)


# Define paths
input_image_path = 'down.jpg'
output_image_path = 'downout2.jpg'
model_path = 'color.onnx'

# Perform colorization
colorize_image(input_image_path, output_image_path, model_path)


python testonnx.py

看看效果

在这里插入图片描述
嗯，模型没有问题，下面开始用c++推理

C++ 推理

未完待续

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：/a/568545.html

如若内容造成侵权/违法违规/事实不符，请联系我们进行投诉反馈qq邮箱809451989@qq.com，一经查实，立即删除！

相关文章

代码随想录算法训练营DAY32|C++贪心算法Part.2|122.买卖股票的最佳时机II、55.跳跃游戏、45.跳跃游戏II

代码随想录算法训练营DAY32|C++贪心算法Part.2|122.买卖股票的最佳时机II、55.跳跃游戏、45.跳跃游戏II

文章目录 122.买卖股票的最佳时机II思路CPP代码 55.跳跃游戏思路CPP代码 45.跳跃游戏II思路方法一代码改善 CPP代码 122.买卖股票的最佳时机II 力扣题目链接 文章讲解&#xff1a;122.买卖股票的最佳时机II 视频讲解&#xff1a; 状态&#xff1a;本题可以用动态规划&#xff0…
阅读更多...
Linux进程地址空间及其页表

Linux进程地址空间及其页表

文章目录 一、引言二、Linux进程地址空间的概念1、进程地址空间定义2、进程地址空间的组成3、进程地址空间与物理内存的关系 三、页表与内存映射1、页表的定义及作用2、页表的缺页中断 三、进程的写时拷贝 一、引言 在Linux中&#xff0c;进程管理是其核心功能之一&#xff0c…
阅读更多...
OpenCV如何实现拉普拉斯算子的离散模拟

OpenCV如何实现拉普拉斯算子的离散模拟

返回:OpenCV系列文章目录&#xff08;持续更新中......&#xff09; 上一篇&#xff1a;OpenCV的Sobel 衍生品 下一篇 &#xff1a;OpenCV 如何实现边缘检测器 目标 在本教程中&#xff0c;您将学习如何&#xff1a; 使用 OpenCV 函数 Laplacian&#xff08;&#xff09; 实…
阅读更多...
Docker基本管理和虚拟化

Docker基本管理和虚拟化

一、docker的发展历史 https://www.cnblogs.com/rongba/articles/14782624.htmlhttps://www.cnblogs.com/rongba/articles/14782624.html 二、docker的概述 Docker是一个开源的应用容器引擎&#xff0c;基于go语言开发并遵循了apache2.0协议开源。 Docker是在Linux容器里运行…
阅读更多...
Amazon云计算AWS之[2]弹性计算云EC2

Amazon云计算AWS之[2]弹性计算云EC2

文章目录 说明EC2基本架构Amazon机器映象&#xff08;AMI&#xff09;实例&#xff08;Instance&#xff09;弹性块存储&#xff08;EBS&#xff09; EC2关键技术地理区域和可用区域EC2通信机制弹性负载均衡监控服务自动缩放服务管理控制台 EC2安全及容错机制EC2弹性IP地址 说明…
阅读更多...
开曼群岛：Web3企业的乐园

开曼群岛：Web3企业的乐园

开曼群岛&#xff1a;Web3企业的理想之地 开曼群岛&#xff0c;在数字革命中大放异彩。近年来&#xff0c;该地区成立的Web3企业数量显著增加&#xff0c;如果保持目前的发展速度&#xff0c;并持续优化立法&#xff0c;那么扩展的速度将无可限量。本文将探讨推动这一增长的关…
阅读更多...
使用Pycharm运行spark实例时没有pyspark包（ModuleNotFoundError: No module named ‘py4j‘）

使用Pycharm运行spark实例时没有pyspark包（ModuleNotFoundError: No module named ‘py4j‘）

一、问题描述 在安装并配置pyspark&#xff0c;下载并打开Pycharm&#xff08;专业版&#xff09;后进行spark实例操作&#xff08;笔者以统计文件中的行数为例&#xff09;时&#xff0c;运行程序后提示ModuleNotFoundError: No module named py4j&#xff1a; 二、解决办法 …
阅读更多...
小程序线多点路图绘制

小程序线多点路图绘制

需求 当接口返回一连串地图坐标&#xff0c;需要根据这些坐标串联起来&#xff0c;形成一个线路图&#xff08;本次使用步行导航线路图&#xff09;。 思路 首先优先想到使用小程序Map组件的polyline属性去进行展示。但是我们发现直接使用该属性进行坐标绘制画出来的数据都是…
阅读更多...
恶补《操作系统》2_1——王道学习笔记

恶补《操作系统》2_1——王道学习笔记

2操作系统-进程 2.1_1 进程的定义、组成、组织方式、特征 组成&#xff1a;PCB&#xff08;进程存在唯一的标志&#xff09;&#xff0c;程序段&#xff0c;数据段 组织方式&#xff1a;链接方式&#xff0c;指针指向不同的队列&#xff1b;索引方式&#xff0c;索引表 特征…
阅读更多...
fnm：Rust开发的高效Node版本管理工具

fnm：Rust开发的高效Node版本管理工具

简介 fnm 是一个基于 Rust 开发的 Node 版本管理工具&#xff0c;它的目标是提供一个快速、简单且可靠的方式来管理 Node.js 的不同版本。同时&#xff0c;它是跨平台的&#xff0c;支持 macOS、Linux、Windows。&#x1f680; Fast and simple Node.js version manager, buil…
阅读更多...
ISP比普通的静态代理相比有什么优势？

ISP比普通的静态代理相比有什么优势？

ISP&#xff08;Internet Service Provider&#xff09;&#xff0c;即互联网服务提供商&#xff0c;是向广大用户综合提供互联网接入业务、信息业务、增值业务的电信运营商。而静态代理则是一个固定不变的代理IP地址&#xff0c;具有稳定性强、兼容性好和管理方便等特点。当我…
阅读更多...
QT中基于TCP的网络通信

QT中基于TCP的网络通信

QT中基于TCP的网络通信 QTcpServer公共成员函数信号 QTcpSocket公共成员函数信号 通信流程服务器端通信流程代码 客户端通信流程代码 使用Qt提供的类进行基于TCP的套接字通信需要用到两个类&#xff1a; QTcpServer&#xff1a;服务器类&#xff0c;用于监听客户端连接以及和客…
阅读更多...
再谈C语言——理解指针（二）

再谈C语言——理解指针（二）

指针变量类型的意义 指针变量的⼤⼩和类型⽆关&#xff0c;只要是指针变量&#xff0c;在同⼀个平台下&#xff0c;⼤⼩都是⼀样的&#xff0c;为什么还要有各种各样的指针类型呢&#xff1f; 其实指针类型是有特殊意义的&#xff0c;我们接下来继续学习。 指针的解引⽤ 对⽐…
阅读更多...
Linux下的UDEV机制/守护进程

Linux下的UDEV机制/守护进程

一. Udev机制概念引入 ( 需要在 etc/udev/rules.d/ 下创建设备的相关规则&#xff0c;不然有可能udev机制生成的设备文件不具备可读可写的权限&#xff0c;adb无法成功通过该设备文件访问设备 ) a. 创建文件夹 sudo vim Xiaomi-audroid.rules b. 添加规则 …
阅读更多...
Laravel 6 - 第十四章 响应

Laravel 6 - 第十四章 响应

​ 文章目录 Laravel 6 - 第一章 简介 Laravel 6 - 第二章 项目搭建 Laravel 6 - 第三章 文件夹结构 Laravel 6 - 第四章 生命周期 Laravel 6 - 第五章 控制反转和依赖注入 Laravel 6 - 第六章 服务容器 Laravel 6 - 第七章 服务提供者 Laravel 6 - 第八章 门面 Laravel 6 - …
阅读更多...
(ICML-2021)从自然语言监督中学习可迁移的视觉模型

(ICML-2021)从自然语言监督中学习可迁移的视觉模型

从自然语言监督中学习可迁移的视觉模型 Title&#xff1a;Learning Transferable Visual Models From Natural Language Supervision paper是OpenAI发表在ICML 21的工作 paper链接 Abstract SOTA计算机视觉系统经过训练可以预测一组固定的预定目标类别。这种受限的监督形式限制…
阅读更多...
LabVIEW与Modbus协议的多点温度监控系统

LabVIEW与Modbus协议的多点温度监控系统

LabVIEW与Modbus协议的多点温度监控系统 随着工业自动化和智能化水平的不断提升&#xff0c;对于现场监控技术的需求日益增长。开发了一种基于LabVIEW与Modbus协议的多点温度监控系统&#xff0c;实现高效、准确的温度数据采集、处理和显示&#xff0c;以及数据存储功能&#…
阅读更多...
c++设计模式之桥接模式（拼接组合）

c++设计模式之桥接模式（拼接组合）

桥接模式&#xff1a;就是进行拼接组装 应用举例&#xff1a; 1.定义了形状&#xff0c;抽象形状接口&#xff0c;圆&#xff0c;矩形 2.定义了颜色&#xff0c;抽象颜色接口&#xff0c;红色&#xff0c;蓝色 3&#xff0c;怎么桥接&#xff0c;抽象具体形状和具体颜色的组合…
阅读更多...
酷开科技逐步为用户构建健全的智慧家庭生活场景

酷开科技逐步为用户构建健全的智慧家庭生活场景

大规模与精细化人群技术则是通过大量的计算能力和精细化的运营能力&#xff0c;建立用户专属数据储存区域&#xff0c;使得用户在使用不同电视的观影偏好和兴趣能够能够得以延续。 不拘泥于自有品牌终端数量&#xff0c;酷开系统除了集成在创维电视上&#xff0c;还服务于飞利…
阅读更多...
大模型的实践应用22-谷歌Gemma AI大模型的架构原理,以及Gemma模型的部署安装本地教程

大模型的实践应用22-谷歌Gemma AI大模型的架构原理,以及Gemma模型的部署安装本地教程

大家好，我是微学AI，今天给大家介绍一下大模型的实践应用22-谷歌Gemma AI大模型的架构原理,以及Gemma模型的部署安装本地教程。谷歌Gemma AI大模型是由Google AI团队开发并开源。Gemma模型采用Transformer编码器-解码器架构，并加入了一些改进，例如使用稀疏注意力机制来提高推…
阅读更多...
最新文章

Copyright @ 2022~2023