koba-jon / pytorch_cpp

Deep Learning sample programs using PyTorch in C++
MIT License
245 stars 52 forks source link

Heatmap from anomaly detection results. #29

Closed sFaster closed 1 year ago

sFaster commented 1 year ago

Hi. Thanks for you codes.

I need to know how heatmap make from anomaly result image and origin image.

Thanks.

koba-jon commented 1 year ago

Hi! For origin image image and generated image output, the heatmap image is created by the following:

torch::Tensor image, output, sub;
cv::Mat sub_mat, heatmap;

sub = torch::abs(image - output).squeeze(/*dim=*/0).mean(/*dim=*/0).unsqueeze(/*dim=*/-1);  // sub{1, C, H, W} ===> {C, H, W} ===> {H, W} ===> {H, W, 1}
sub = sub / 2.0 * 255.0;  // sub[0, 2] ===> [0, 255]
sub = sub.to(torch::kCPU);  // CUDA ===> CPU
sub = sub.to(torch::kUInt8);  // float ===> unsigned char
sub_mat = cv::Mat(cv::Size(sub.size(1), sub.size(0)), CV_8UC1, sub.data_ptr<unsigned char>());  // torch::Tensor ===> cv::Mat
cv::applyColorMap(sub_mat, heatmap, cv::COLORMAP_JET);  // make heatmap
cv::imwrite("save_image_path.png", heatmap);

Thank you.

koba-jon commented 1 year ago

And, I released the new code for creating the heatmap as following: https://github.com/koba-jon/pytorch_cpp/tree/support/v2.0.1

Thank you.

sFaster commented 1 year ago

Thank you. :)