[PyTorch & LibTorch] OpenCV Image to Normalized Tensor C++ 코드

1. 이미지를 Load 한다.

 

2. 이미지를 Resize 한다.

 

3. (모델이 요구하는 인풋 텐서의 컬러 포맷이 RGB인 경우) 이미지의 컬러 포맷을 RGB로 변환한다. 

*OpenCV는 기본적으로 BGR 포맷으로 이미지를 처리한다.

 

4. 이미지를 텐서(torch::Tensor 형)로 변환한다.

 

5. 모델이 요구하는 인풋 텐서의 Shape에 맞도록 텐서를 Reshape한다. (본 예제에서는 NCHW 포맷을 따름)

 

6. 텐서를 정규화한다.

 

7. (GPU를 사용하는 경우)  텐서를 GPU 메모리로 업로드한다.

 

cv::Mat img_bgr_u8 = cv::imread("your_image.jpg",cv::IMREAD_COLOR);
cv::Mat img_rgb_u8;

//resize
cv::resize(img_bgr_u8, img_bgr_u8, cv::Size(416, 416)); #size is dependent on your network

//bgr 2 rgb
cv::cvtColor(img_bgr_u8, img_rgb_u8, cv::COLOR_BGR2RGB);

//to tensor
torch::Tensor input_tensor = torch::from_blob(img_rgb_u8.data, {img_rgb_u8.rows, img_rgb_u8.cols, 3}, torch::kByte);

//transpose
input_tensor = input_tensor.permute({2, 0, 1});
input_tensor = input_tensor.toType(at::kFloat);

//normalize, this example assusmes that input is normalized by Imagenet statistics.
const float mean[3] = {0.485, 0.456, 0.406};
const float std[3] = {0.229, 0.224, 0.225};

input_tensor.div_(255.);
for(int ch=0; ch < 3; ch++)
{
	input_tensor[ch].sub_(mean[ch]).div_(std[ch]);
}
input_tensor.unsqueeze_(0);

//if you use gpu for inference
input_tensor = input_tensor.to(torch::kCUDA);