选择的模型是Zero-DCE_extension，虽然是几年前的论文但它的原理简洁且合理，泛化性比一些基于成对图像训练的监督模型要好很多。它没有基于图像映射实现光照增强，而是基于对图像中像素的估计，通过模型拟合的高阶曲线对原图进行光照增强。可以去看看这篇解析博客讲的非常详细，通过ps中的提亮曲线可以很轻松理解本文idea的可行性。
提亮曲线
既然本文的重点在于提亮曲线，那就设计到两部分内容：
- 模型的结构对参数估计的影响
- 提亮曲线的函数设计
模型结构设计
从模型结构来看其实非常简单，就是由多层深度可分离卷积与残差结构组成，最后输出图像中各个位置在三个不同维度上的估计值，再运用设计的提亮曲线增强原始图像最终输出增强后图像。也正是因为模型简单，同时基本用的也都是3*3的小尺度卷积核，所以对于局部像素估计值的精度可能存在偏差。
提亮曲线设计
文中提出提亮曲线必须符合下面三个要求：
- 函数的值域范围应该是[0,1]之间的，且x=0时函数值为0，x=1时函数值为1；
- 函数具有单调性；
- 函数为了符合反向传播，所以必须是可导的；
本文提出的提亮曲线设计为f(x,alpha) = x + alpha*x*(1-x)，其中x表示图像中某个位置的像素值，alpha为模型估计出的参数值。然后多次迭代这个曲线，即g(x,alpha)=f(f(...f(x,alpha)))。这个函数是否就是最优解呢？把相关的要求问一下KIMI，得到的答案如下：
这个改进之后的曲线就不再那么平滑，在多次迭代后的梯度会更大。优点是可能模型会更快收敛，缺点则是梯度过大容易震荡，同时损失函数中的权重参数也更加难设置。
模型导出与推理
模型导出可以直接参考下面的代码
import torch
import onnx
from model import infer
model = infer()
model.eval()
ckpts = torch.load('./new_ckpts/Epoch129_111.091.pth')
model.load_state_dict(ckpts)
x = torch.rand(1,3,512,512)
input_name = "images"
output_name = "output"
torch.onnx.export(model,x,"zero_DCE_plus.onnx",opset_version=13,
input_names=[input_name],output_names=[output_name],
dynamic_axes={
input_name:{0:'batch_size'},
output_name:{0:'batch_size'},
}
)
得到onnx文件后转为TensorRT的engine文件
trtexec --onnx=zero_DCE_plus.onnx --saveEngine=DCE_plus_fp16.engine --fp16
推理部分
TensorRT采用10.X版本，相关engine反序列化及文件解析可以参考下面代码
this->num_bindings = this->engine->getNbIOTensors();
for (int i = 0; i < this->num_bindings; ++i) {
Binding binding;
nvinfer1::Dims dims;
auto name = this->engine->getIOTensorName(i);
nvinfer1::DataType dtype = this->engine->getTensorDataType(name);
binding.name = name;
binding.dsize = type_to_size(dtype);
dims = this->engine->getTensorShape(name);
auto ioType = this->engine->getTensorIOMode(name);
if(ioType == nvinfer1::TensorIOMode::kINPUT) {
this->num_inputs += 1;
dims = this->engine->getProfileShape(name,0,nvinfer1::OptProfileSelector::kMAX);
binding.size = get_size_by_dims(dims);
binding.dims = dims;
this->input_bindings.push_back(binding);
this->context->setInputShape(name, dims);
}else {
binding.size = get_size_by_dims(dims);
binding.dims = dims;
this->output_bindings.push_back(binding);
this->num_outputs += 1;
}
}
和8.X版本略有不同，但是差异不大。
为了进一步降低推理耗时，自己写前处理（blobfromimage）来代替opencv的函数，会有一定提升。
auto st = std::chrono::high_resolution_clock::now();
cv::dnn::blobFromImage(nchw, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
auto end = std::chrono::high_resolution_clock::now();
auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - st);
std::cout <<"opencv blob cost: "<< duration.count() << "ms \n";
writeFloatArrayToFile(out.ptr<float>(0), nchw.total() * 3, "../blob.csv");
st = std::chrono::high_resolution_clock::now();
cv::cvtColor(nchw, nchw, cv::COLOR_BGR2RGB);
float* data = new float[nchw.total() * 3];
blobFromImage(nchw,data); //speed up
end = std::chrono::high_resolution_clock::now();
duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - st);
std::cout <<"our blob cost: "<< duration.count() << "ms \n";
writeFloatArrayToFile(data, nchw.total() * 3, "../myblob.csv");
再来看看写入文件的数据是否一致
通过脚本检查，处理后的数据和opencv函数输出完全相同，但是耗时从8ms降至1ms还是在cpu上，如果还需要优化可以写成cuda的预处理kernel，这里暂时不做进一步优化了。后处理部分与前处理类似，直接将模型输出的值转为Mat的data
const int step = ow*oh;
cv::Mat out = cv::Mat::zeros(ow,oh,CV_8UC3);
unsigned char* data = out.data;
// copy data to img && RGB -> BGR
for(int i=0;i<oh;++i) {
int index = i*ow;
for(int j=0;j<ow;++j) {
int tmp = (index+j) * 3;
data[tmp] = std::min(255,int(ptr[index + j + 2 * step]*255));
data[tmp + 1] = std::min(255,int(ptr[index + j + step]*255));
data[tmp + 2] = std::min(255,int(ptr[index + j]*255));
}
}
cv::resize(out,image,cv::Size(dw,dh));
运行测试
遍历文件夹对图像进行增亮处理，运行耗时如下：
每张图片总体耗时大约为4~5ms，也就是200FPS以上。
下面展示一些结果图像
一些问题
- 模型结构是否可以改进为基于patch的局部像素块进行参数估计，更加贴近图像亮度分布
- 提亮函数设计还可以考虑更多可能性
- 优化提亮后图像质量，原始代码的损失中并没有对提亮后图像进行总变分损失计算
- 如需进一步提升推理速度，可以考虑先下采样再上采样，同原始代码中进行相似操作
本文的相关代码可以在DCE_CPP中查看
