阿里开源黑白图片上色算法DDColor的部署与测试并将模型转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 Fclass ImageColorizationPipeline(object):def __init__(self, model_path, input_size=256, model_size='large'):self.input_size = input_sizeif 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 = 256img = (img / 255.0).astype(np.float32)orig_l = cv2.cvtColor(img, cv2