├── README.md ├── main.cpp ├── main.py └── testimgs ├── indoor ├── 1403_7.png ├── 1413_2.png ├── 1420_10.png ├── 1425_10.png ├── 1430_10.png ├── 1434_10.png └── 1444_10.png └── outdoor ├── 0011_0.95_0.16.jpg ├── 0036_1_0.08.jpg ├── 0046_0.9_0.16.jpg ├── 0061_0.8_0.2.jpg ├── 0076_1_0.16.jpg ├── 0108_1_0.2.jpg ├── 0112_0.8_0.2.jpg ├── 0132_1_0.08.jpg ├── 0137_0.8_0.2.jpg ├── 0143_1_0.2.jpg ├── 0160_0.9_0.2.jpg ├── 0166_0.8_0.16.jpg ├── 0226_0.8_0.2.jpg ├── 0273_0.85_0.2.jpg ├── 0300_1_0.2.jpg ├── 0335_1_0.2.jpg ├── 0393_1_0.2.jpg ├── 0406_1_0.2.jpg └── 1038_1_0.16.jpg /README.md: -------------------------------------------------------------------------------- 1 | 模型分为indoor室内和outdoor室外这两套场景的，起初我是用opencv-dnn做推理，可是输出结果图是全黑的。 2 | 换做使用onnxruntime做推理，输出结果图是合理的。 3 | 4 | onnx文件在百度云盘，链接：https://pan.baidu.com/s/10O6IQ7bEMIPrM4Mzt349cw 5 | 提取码：wwh4 6 | 7 | 训练源码在 https://github.com/YuZheng9/C2PNet 8 | 它是中了cvpr2023文章 9 | -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | //#include ///如果使用cuda加速，需要取消注释 7 | #include 8 | 9 | using namespace cv; 10 | using namespace std; 11 | using namespace Ort; 12 | 13 | class C2PNet 14 | { 15 | public: 16 | C2PNet(string modelpath); 17 | Mat predict(Mat srcimg); 18 | private: 19 | void preprocess(Mat img); 20 | vector input_image; 21 | int inpWidth; 22 | int inpHeight; 23 | bool dynamic_shape; 24 | 25 | Env env = Env(ORT_LOGGING_LEVEL_ERROR, "Image Dehazing"); 26 | Ort::Session *ort_session = nullptr; 27 | SessionOptions sessionOptions = SessionOptions(); 28 | vector input_names; 29 | vector output_names; 30 | vector> input_node_dims; // >=1 outputs 31 | vector> output_node_dims; // >=1 outputs 32 | Ort::MemoryInfo memory_info_handler = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); 33 | }; 34 | 35 | C2PNet::C2PNet(string model_path) 36 | { 37 | std::wstring widestr = std::wstring(model_path.begin(), model_path.end()); ////windows写法 38 | ///OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0); ///如果使用cuda加速，需要取消注释 39 | 40 | sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC); 41 | ort_session = new Session(env, widestr.c_str(), sessionOptions); ////windows写法 42 | ////ort_session = new Session(env, model_path.c_str(), sessionOptions); ////linux写法 43 | 44 | size_t numInputNodes = ort_session->GetInputCount(); 45 | size_t numOutputNodes = ort_session->GetOutputCount(); 46 | AllocatorWithDefaultOptions allocator; 47 | for (int i = 0; i < numInputNodes; i++) 48 | { 49 | input_names.push_back(ort_session->GetInputName(i, allocator)); 50 | Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i); 51 | auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo(); 52 | auto input_dims = input_tensor_info.GetShape(); 53 | input_node_dims.push_back(input_dims); 54 | 55 | } 56 | for (int i = 0; i < numOutputNodes; i++) 57 | { 58 | output_names.push_back(ort_session->GetOutputName(i, allocator)); 59 | Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i); 60 | auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo(); 61 | auto output_dims = output_tensor_info.GetShape(); 62 | output_node_dims.push_back(output_dims); 63 | } 64 | 65 | this->inpHeight = input_node_dims[0][2]; 66 | this->inpWidth = input_node_dims[0][3]; 67 | if (this->inpHeight == -1 || this->inpWidth == -1) 68 | { 69 | this->dynamic_shape = true; 70 | } 71 | else 72 | { 73 | this->dynamic_shape = false; 74 | } 75 | } 76 | 77 | void C2PNet::preprocess(Mat img) 78 | { 79 | Mat rgbimg; 80 | cvtColor(img, rgbimg, COLOR_BGR2RGB); 81 | if (!dynamic_shape) 82 | { 83 | resize(rgbimg, rgbimg, cv::Size(this->inpWidth, this->inpHeight)); 84 | } 85 | else 86 | { 87 | this->inpHeight = rgbimg.rows; 88 | this->inpWidth = rgbimg.cols; 89 | } 90 | 91 | vector rgbChannels(3); 92 | split(rgbimg, rgbChannels); 93 | for (int c = 0; c < 3; c++) 94 | { 95 | rgbChannels[c].convertTo(rgbChannels[c], CV_32FC1, 1 / 255.0); 96 | } 97 | 98 | const int image_area = this->inpHeight * this->inpWidth; 99 | this->input_image.resize(3 * image_area); 100 | size_t single_chn_size = image_area * sizeof(float); 101 | memcpy(this->input_image.data(), (float*)rgbChannels[0].data, single_chn_size); 102 | memcpy(this->input_image.data() + image_area, (float*)rgbChannels[1].data, single_chn_size); 103 | memcpy(this->input_image.data() + image_area * 2, (float*)rgbChannels[2].data, single_chn_size); 104 | } 105 | 106 | Mat C2PNet::predict(Mat srcimg) 107 | { 108 | const int srch = srcimg.rows; 109 | const int srcw = srcimg.cols; 110 | this->preprocess(srcimg); 111 | 112 | std::vector input_img_shape = { 1, 3, this->inpHeight, this->inpWidth }; 113 | Value input_tensor_ = Value::CreateTensor(memory_info_handler, this->input_image.data(), this->input_image.size(), input_img_shape.data(), input_img_shape.size()); 114 | 115 | Ort::RunOptions runOptions; 116 | vector ort_outputs = this->ort_session->Run(runOptions, this->input_names.data(), &input_tensor_, 1, this->output_names.data(), output_names.size()); 117 | 118 | float* pdata = ort_outputs[0].GetTensorMutableData(); 119 | std::vector outs_shape = ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape(); 120 | const int out_h = outs_shape[2]; 121 | const int out_w = outs_shape[3]; 122 | const int channel_step = out_h * out_w; 123 | Mat rmat(out_h, out_w, CV_32FC1, pdata); 124 | Mat gmat(out_h, out_w, CV_32FC1, pdata + channel_step); 125 | Mat bmat(out_h, out_w, CV_32FC1, pdata + 2 * channel_step); 126 | rmat *= 255.f; 127 | gmat *= 255.f; 128 | bmat *= 255.f; 129 | ///output_image = 等价np.clip(output_image, 0, 255) 130 | rmat.setTo(0, rmat < 0); 131 | rmat.setTo(255, rmat > 255); 132 | gmat.setTo(0, gmat < 0); 133 | gmat.setTo(255, gmat > 255); 134 | bmat.setTo(0, bmat < 0); 135 | bmat.setTo(255, bmat > 255); 136 | 137 | vector channel_mats(3); 138 | channel_mats[0] = bmat; 139 | channel_mats[1] = gmat; 140 | channel_mats[2] = rmat; 141 | 142 | Mat dstimg; 143 | merge(channel_mats, dstimg); 144 | dstimg.convertTo(dstimg, CV_8UC3); 145 | resize(dstimg, dstimg, cv::Size(srcw, srch)); 146 | return dstimg; 147 | } 148 | 149 | int main() 150 | { 151 | C2PNet mynet("weights/c2pnet_indoor_HxW.onnx"); 152 | 153 | string imgpath = "testimgs/indoor/1420_10.png"; 154 | Mat srcimg = imread(imgpath); 155 | Mat dstimg = mynet.predict(srcimg); 156 | 157 | namedWindow("srcimg", WINDOW_NORMAL); 158 | imshow("srcimg", srcimg); 159 | static const string kWinName = "Deep learning Image Dehaze use ONNXRuntime"; 160 | namedWindow(kWinName, WINDOW_NORMAL); 161 | imshow(kWinName, dstimg); 162 | waitKey(0); 163 | destroyAllWindows(); 164 | 165 | return 0; 166 | } 167 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import cv2 3 | import onnxruntime 4 | import numpy as np 5 | 6 | 7 | class C2PNet: 8 | def __init__(self, modelpath): 9 | # Initialize model 10 | self.onnx_session = onnxruntime.InferenceSession(modelpath) 11 | self.input_name = self.onnx_session.get_inputs()[0].name 12 | _, _, self.input_height, self.input_width = self.onnx_session.get_inputs()[0].shape 13 | 14 | def detect(self, image): 15 | input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 16 | if isinstance(self.input_height ,int) and isinstance(self.input_width, int): 17 | input_image = cv2.resize(input_image, (self.input_width, self.input_height)) ###固定输入分辨率, HXW.onnx文件是动态输入分辨率的 18 | input_image = input_image.astype(np.float32) / 255.0 19 | input_image = input_image.transpose(2, 0, 1) 20 | input_image = np.expand_dims(input_image, axis=0) 21 | 22 | result = self.onnx_session.run(None, {self.input_name: input_image}) ###opencv-dnn推理时,结果图全黑 23 | 24 | # Post process:squeeze, RGB->BGR, Transpose, uint8 cast 25 | output_image = np.squeeze(result[0]) 26 | output_image = output_image.transpose(1, 2, 0) 27 | output_image = output_image * 255 28 | output_image = np.clip(output_image, 0, 255) 29 | output_image = output_image.astype(np.uint8) 30 | output_image = cv2.cvtColor(output_image.astype(np.uint8), cv2.COLOR_RGB2BGR) 31 | output_image = cv2.resize(output_image, (image.shape[1], image.shape[0])) 32 | return output_image 33 | 34 | 35 | if __name__ == '__main__': 36 | parser = argparse.ArgumentParser() 37 | parser.add_argument('--imgpath', type=str, 38 | default='testimgs/outdoor/0143_1_0.2.jpg', help="image path") 39 | parser.add_argument('--modelpath', type=str, 40 | default='weights/c2pnet_outdoor_360x640.onnx', help="onnx path") 41 | args = parser.parse_args() 42 | 43 | mynet = C2PNet(args.modelpath) 44 | srcimg = cv2.imread(args.imgpath) 45 | 46 | dstimg = mynet.detect(srcimg) 47 | 48 | if srcimg.shape[0] > srcimg.shape[1]: 49 | boundimg = np.zeros((10, srcimg.shape[1], 3), dtype=srcimg.dtype)+255 ###中间分开原图和结果 50 | combined_img = np.vstack([srcimg, boundimg, dstimg]) 51 | else: 52 | boundimg = np.zeros((srcimg.shape[0], 10, 3), dtype=srcimg.dtype)+255 53 | combined_img = np.hstack([srcimg, boundimg, dstimg]) 54 | cv2.imwrite('result.jpg', combined_img) 55 | winName = 'Deep learning Image Dehaze use onnxruntime' 56 | cv2.namedWindow(winName, 0) 57 | cv2.imshow(winName, combined_img) ###原图和结果图也可以分开窗口显示 58 | cv2.waitKey(0) 59 | cv2.destroyAllWindows() 60 | -------------------------------------------------------------------------------- /testimgs/indoor/1403_7.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1403_7.png -------------------------------------------------------------------------------- /testimgs/indoor/1413_2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1413_2.png -------------------------------------------------------------------------------- /testimgs/indoor/1420_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1420_10.png -------------------------------------------------------------------------------- /testimgs/indoor/1425_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1425_10.png -------------------------------------------------------------------------------- /testimgs/indoor/1430_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1430_10.png -------------------------------------------------------------------------------- /testimgs/indoor/1434_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1434_10.png -------------------------------------------------------------------------------- /testimgs/indoor/1444_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/indoor/1444_10.png -------------------------------------------------------------------------------- /testimgs/outdoor/0011_0.95_0.16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0011_0.95_0.16.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0036_1_0.08.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0036_1_0.08.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0046_0.9_0.16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0046_0.9_0.16.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0061_0.8_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0061_0.8_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0076_1_0.16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0076_1_0.16.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0108_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0108_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0112_0.8_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0112_0.8_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0132_1_0.08.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0132_1_0.08.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0137_0.8_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0137_0.8_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0143_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0143_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0160_0.9_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0160_0.9_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0166_0.8_0.16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0166_0.8_0.16.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0226_0.8_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0226_0.8_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0273_0.85_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0273_0.85_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0300_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0300_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0335_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0335_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0393_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0393_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/0406_1_0.2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/0406_1_0.2.jpg -------------------------------------------------------------------------------- /testimgs/outdoor/1038_1_0.16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/C2PNet-onnxrun/0207db5b11f64c7fefd1869aee89a4b1ab05c2e5/testimgs/outdoor/1038_1_0.16.jpg --------------------------------------------------------------------------------