基于深度学习神经网络的AI图片上色DDcolor系统源码

第一步：DDcolor介绍

DDColor 是最新的 SOTA 图像上色算法，能够对输入的黑白图像生成自然生动的彩色结果，使用 UNet 结构的骨干网络和图像解码器分别实现图像特征提取和特征图上采样，并利用 Transformer 结构的颜色解码器完成基于视觉语义的颜色查询，最终聚合输出彩色通道预测结果。

它甚至可以对动漫游戏中的风景进行着色/重新着色，将您的动画风景转变为逼真的现实生活风格！（图片来源：原神）

第二步：DDcolor网络结构

算法整体流程如下图，使用 UNet 结构的骨干网络和图像解码器分别实现图像特征提取和特征图上采样，并利用 Transformer 结构的颜色解码器完成基于视觉语义的颜色查询，最终聚合输出彩色通道预测结果。

第三步：模型代码展示

import os
import torch
from collections import OrderedDict
from os import path as osp
from tqdm import tqdm
import numpy as npfrom basicsr.archs import build_network
from basicsr.losses import build_loss
from basicsr.metrics import calculate_metric
from basicsr.utils import get_root_logger, imwrite, tensor2img
from basicsr.utils.img_util import tensor_lab2rgb
from basicsr.utils.dist_util import master_only
from basicsr.utils.registry import MODEL_REGISTRY
from .base_model import BaseModel
from basicsr.metrics.custom_fid import INCEPTION_V3_FID, get_activations, calculate_activation_statistics, calculate_frechet_distance
from basicsr.utils.color_enhance import color_enhacne_blend@MODEL_REGISTRY.register()
class ColorModel(BaseModel):"""Colorization model for single image colorization."""def __init__(self, opt):super(ColorModel, self).__init__(opt)# define network net_gself.net_g = build_network(opt['network_g'])self.net_g = self.model_to_device(self.net_g)self.print_network(self.net_g)# load pretrained model for net_gload_path = self.opt['path'].get('pretrain_network_g', None)if load_path is not None:param_key = self.opt['path'].get('param_key_g', 'params')self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key)if self.is_train:self.init_training_settings()def init_training_settings(self):train_opt = self.opt['train']self.ema_decay = train_opt.get('ema_decay', 0)if self.ema_decay > 0:logger = get_root_logger()logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')# define network net_g with Exponential Moving Average (EMA)# net_g_ema is used only for testing on one GPU and saving# There is no need to wrap with DistributedDataParallelself.net_g_ema = build_network(self.opt['network_g']).to(self.device)# load pretrained modelload_path = self.opt['path'].get('pretrain_network_g', None)if load_path is not None:self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')else:self.model_ema(0)  # copy net_g weightself.net_g_ema.eval()# define network net_dself.net_d = build_network(self.opt['network_d'])self.net_d = self.model_to_device(self.net_d)self.print_network(self.net_d)# load pretrained model for net_dload_path = self.opt['path'].get('pretrain_network_d', None)if load_path is not None:param_key = self.opt['path'].get('param_key_d', 'params')self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True), param_key)self.net_g.train()self.net_d.train()# define lossesif train_opt.get('pixel_opt'):self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)else:self.cri_pix = Noneif train_opt.get('perceptual_opt'):self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)else:self.cri_perceptual = Noneif train_opt.get('gan_opt'):self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)else:self.cri_gan = Noneif self.cri_pix is None and self.cri_perceptual is None:raise ValueError('Both pixel and perceptual losses are None.')if train_opt.get('colorfulness_opt'):self.cri_colorfulness = build_loss(train_opt['colorfulness_opt']).to(self.device)else:self.cri_colorfulness = None# set up optimizers and schedulersself.setup_optimizers()self.setup_schedulers()# set real dataset cache for fid metric computingself.real_mu, self.real_sigma = None, Noneif self.opt['val'].get('metrics') is not None and self.opt['val']['metrics'].get('fid') is not None:self._prepare_inception_model_fid()def setup_optimizers(self):train_opt = self.opt['train']# optim_params_g = []# for k, v in self.net_g.named_parameters():#     if v.requires_grad:#         optim_params_g.append(v)#     else:#         logger = get_root_logger()#         logger.warning(f'Params {k} will not be optimized.')optim_params_g = self.net_g.parameters()# optimizer goptim_type = train_opt['optim_g'].pop('type')self.optimizer_g = self.get_optimizer(optim_type, optim_params_g, **train_opt['optim_g'])self.optimizers.append(self.optimizer_g)# optimizer doptim_type = train_opt['optim_d'].pop('type')self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])self.optimizers.append(self.optimizer_d)def feed_data(self, data):self.lq = data['lq'].to(self.device)self.lq_rgb = tensor_lab2rgb(torch.cat([self.lq, torch.zeros_like(self.lq), torch.zeros_like(self.lq)], dim=1))if 'gt' in data:self.gt = data['gt'].to(self.device)self.gt_lab = torch.cat([self.lq, self.gt], dim=1)self.gt_rgb = tensor_lab2rgb(self.gt_lab)if self.opt['train'].get('color_enhance', False):for i in range(self.gt_rgb.shape[0]):self.gt_rgb[i] = color_enhacne_blend(self.gt_rgb[i], factor=self.opt['train'].get('color_enhance_factor'))def optimize_parameters(self, current_iter):# optimize net_gfor p in self.net_d.parameters():p.requires_grad = Falseself.optimizer_g.zero_grad()self.output_ab = self.net_g(self.lq_rgb)self.output_lab = torch.cat([self.lq, self.output_ab], dim=1)self.output_rgb = tensor_lab2rgb(self.output_lab)l_g_total = 0loss_dict = OrderedDict()# pixel lossif self.cri_pix:l_g_pix = self.cri_pix(self.output_ab, self.gt)l_g_total += l_g_pixloss_dict['l_g_pix'] = l_g_pix# perceptual lossif self.cri_perceptual:l_g_percep, l_g_style = self.cri_perceptual(self.output_rgb, self.gt_rgb)if l_g_percep is not None:l_g_total += l_g_perceploss_dict['l_g_percep'] = l_g_percepif l_g_style is not None:l_g_total += l_g_styleloss_dict['l_g_style'] = l_g_style# gan lossif self.cri_gan:fake_g_pred = self.net_d(self.output_rgb)l_g_gan = self.cri_gan(fake_g_pred, target_is_real=True, is_disc=False)l_g_total += l_g_ganloss_dict['l_g_gan'] = l_g_gan# colorfulness lossif self.cri_colorfulness:l_g_color = self.cri_colorfulness(self.output_rgb)l_g_total += l_g_colorloss_dict['l_g_color'] = l_g_colorl_g_total.backward()self.optimizer_g.step()# optimize net_dfor p in self.net_d.parameters():p.requires_grad = Trueself.optimizer_d.zero_grad()real_d_pred = self.net_d(self.gt_rgb)fake_d_pred = self.net_d(self.output_rgb.detach())l_d = self.cri_gan(real_d_pred, target_is_real=True, is_disc=True) + self.cri_gan(fake_d_pred, target_is_real=False, is_disc=True)loss_dict['l_d'] = l_dloss_dict['real_score'] = real_d_pred.detach().mean()loss_dict['fake_score'] = fake_d_pred.detach().mean()l_d.backward()self.optimizer_d.step()self.log_dict = self.reduce_loss_dict(loss_dict)if self.ema_decay > 0:self.model_ema(decay=self.ema_decay)def get_current_visuals(self):out_dict = OrderedDict()out_dict['lq'] = self.lq_rgb.detach().cpu()out_dict['result'] = self.output_rgb.detach().cpu()if self.opt['logger'].get('save_snapshot_verbose', False):  # only for verboseself.output_lab_chroma = torch.cat([torch.ones_like(self.lq) * 50, self.output_ab], dim=1)self.output_rgb_chroma = tensor_lab2rgb(self.output_lab_chroma)out_dict['result_chroma'] = self.output_rgb_chroma.detach().cpu()if hasattr(self, 'gt'):out_dict['gt'] = self.gt_rgb.detach().cpu()if self.opt['logger'].get('save_snapshot_verbose', False):  # only for verboseself.gt_lab_chroma = torch.cat([torch.ones_like(self.lq) * 50, self.gt], dim=1)self.gt_rgb_chroma = tensor_lab2rgb(self.gt_lab_chroma)out_dict['gt_chroma'] = self.gt_rgb_chroma.detach().cpu()return out_dictdef test(self):if hasattr(self, 'net_g_ema'):self.net_g_ema.eval()with torch.no_grad():self.output_ab = self.net_g_ema(self.lq_rgb)self.output_lab = torch.cat([self.lq, self.output_ab], dim=1)self.output_rgb = tensor_lab2rgb(self.output_lab)else:self.net_g.eval()with torch.no_grad():self.output_ab = self.net_g(self.lq_rgb)self.output_lab = torch.cat([self.lq, self.output_ab], dim=1)self.output_rgb = tensor_lab2rgb(self.output_lab)self.net_g.train()def dist_validation(self, dataloader, current_iter, tb_logger, save_img):if self.opt['rank'] == 0:self.nondist_validation(dataloader, current_iter, tb_logger, save_img)def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):dataset_name = dataloader.dataset.opt['name']with_metrics = self.opt['val'].get('metrics') is not Noneuse_pbar = self.opt['val'].get('pbar', False)if with_metrics and not hasattr(self, 'metric_results'):  # only execute in the first runself.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}# initialize the best metric results for each dataset_name (supporting multiple validation datasets)if with_metrics:self._initialize_best_metric_results(dataset_name)# zero self.metric_resultsif with_metrics:self.metric_results = {metric: 0 for metric in self.metric_results}metric_data = dict()if use_pbar:pbar = tqdm(total=len(dataloader), unit='image')if self.opt['val']['metrics'].get('fid') is not None:fake_acts_set, acts_set = [], []for idx, val_data in enumerate(dataloader):# if idx == 100:#     breakimg_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]if hasattr(self, 'gt'):del self.gtself.feed_data(val_data)self.test()visuals = self.get_current_visuals()sr_img = tensor2img([visuals['result']])metric_data['img'] = sr_imgif 'gt' in visuals:gt_img = tensor2img([visuals['gt']])metric_data['img2'] = gt_imgtorch.cuda.empty_cache()if save_img:if self.opt['is_train']:save_dir = osp.join(self.opt['path']['visualization'], img_name)for key in visuals:save_path = os.path.join(save_dir, '{}_{}.png'.format(current_iter, key))img = tensor2img(visuals[key])imwrite(img, save_path)else:if self.opt['val']['suffix']:save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,f'{img_name}_{self.opt["val"]["suffix"]}.png')else:save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,f'{img_name}_{self.opt["name"]}.png')imwrite(sr_img, save_img_path)if with_metrics:# calculate metricsfor name, opt_ in self.opt['val']['metrics'].items():if name == 'fid':pred, gt = visuals['result'].cuda(), visuals['gt'].cuda()fake_act = get_activations(pred, self.inception_model_fid, 1)fake_acts_set.append(fake_act)if self.real_mu is None:real_act = get_activations(gt, self.inception_model_fid, 1)acts_set.append(real_act)else:self.metric_results[name] += calculate_metric(metric_data, opt_)if use_pbar:pbar.update(1)pbar.set_description(f'Test {img_name}')if use_pbar:pbar.close()if with_metrics:if self.opt['val']['metrics'].get('fid') is not None:if self.real_mu is None:acts_set = np.concatenate(acts_set, 0)self.real_mu, self.real_sigma = calculate_activation_statistics(acts_set)fake_acts_set = np.concatenate(fake_acts_set, 0)fake_mu, fake_sigma = calculate_activation_statistics(fake_acts_set)fid_score = calculate_frechet_distance(self.real_mu, self.real_sigma, fake_mu, fake_sigma)self.metric_results['fid'] = fid_scorefor metric in self.metric_results.keys():if metric != 'fid':self.metric_results[metric] /= (idx + 1)# update the best metric resultself._update_best_metric_result(dataset_name, metric, self.metric_results[metric], current_iter)self._log_validation_metric_values(current_iter, dataset_name, tb_logger)def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):log_str = f'Validation {dataset_name}\n'for metric, value in self.metric_results.items():log_str += f'\t # {metric}: {value:.4f}'if hasattr(self, 'best_metric_results'):log_str += (f'\tBest: {self.best_metric_results[dataset_name][metric]["val"]:.4f} @ 'f'{self.best_metric_results[dataset_name][metric]["iter"]} iter')log_str += '\n'logger = get_root_logger()logger.info(log_str)if tb_logger:for metric, value in self.metric_results.items():tb_logger.add_scalar(f'metrics/{dataset_name}/{metric}', value, current_iter)def _prepare_inception_model_fid(self, path='pretrain/inception_v3_google-1a9a5a14.pth'):incep_state_dict = torch.load(path, map_location='cpu')block_idx = INCEPTION_V3_FID.BLOCK_INDEX_BY_DIM[2048]self.inception_model_fid = INCEPTION_V3_FID(incep_state_dict, [block_idx])self.inception_model_fid.cuda()self.inception_model_fid.eval()@master_onlydef save_training_images(self, current_iter):visuals = self.get_current_visuals()save_dir = osp.join(self.opt['root_path'], 'experiments', self.opt['name'], 'training_images_snapshot')os.makedirs(save_dir, exist_ok=True)for key in visuals:save_path = os.path.join(save_dir, '{}_{}.png'.format(current_iter, key))img = tensor2img(visuals[key])imwrite(img, save_path)def save(self, epoch, current_iter):if hasattr(self, 'net_g_ema'):self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])else:self.save_network(self.net_g, 'net_g', current_iter)self.save_network(self.net_d, 'net_d', current_iter)self.save_training_state(epoch, current_iter)

第四步：运行