72 lines
1.7 KiB
C++
72 lines
1.7 KiB
C++
#include <cstring>
|
|
|
|
#include <glog/logging.h>
|
|
#include <sys/stat.h>
|
|
#include <opencv2/core/mat.hpp>
|
|
#include <opencv2/imgcodecs.hpp>
|
|
|
|
#include "PNGDumper.hpp"
|
|
|
|
using namespace fmri;
|
|
using namespace std;
|
|
|
|
static void ensureDir(const string& dir)
|
|
{
|
|
struct stat s;
|
|
if (stat(dir.c_str(), &s) == 0) {
|
|
CHECK(S_ISDIR(s.st_mode)) << dir << " already exists and is not a directory." << endl;
|
|
return;
|
|
}
|
|
|
|
switch (errno) {
|
|
case ENOENT:
|
|
PCHECK(mkdir(dir.c_str(), 0777) == 0) << "Couldn't create directory";
|
|
return;
|
|
|
|
default:
|
|
perror("Unusable dump dir");
|
|
break;
|
|
}
|
|
}
|
|
|
|
PNGDumper::PNGDumper(string_view baseDir) :
|
|
baseDir_(baseDir)
|
|
{
|
|
ensureDir(baseDir_);
|
|
}
|
|
|
|
void PNGDumper::dump(const LayerData &layerData)
|
|
{
|
|
if (layerData.shape().size() == 4) {
|
|
// We have a series of images.
|
|
dumpImageSeries(layerData);
|
|
} else {
|
|
LOG(INFO) << "Unable to dump this type of layer to png.";
|
|
}
|
|
}
|
|
|
|
void PNGDumper::dumpImageSeries(const LayerData &layer)
|
|
{
|
|
const auto& shape = layer.shape();
|
|
const auto images = shape[0], channels = shape[1], height = shape[2], width = shape[3];
|
|
const auto imagePixels = width * height;
|
|
|
|
cv::Mat image(width, height, CV_32FC1);
|
|
|
|
auto data = layer.data();
|
|
|
|
for (int i = 0; i < images; ++i) {
|
|
for (int j = 0; j < channels; ++j) {
|
|
char pathBuf[PATH_MAX];
|
|
std::copy_n(data, imagePixels, image.begin<float>());
|
|
rescale(image.begin<float>(), image.end<float>(), 0, 255);
|
|
std::snprintf(pathBuf, sizeof(pathBuf), "%s/%s-%d-%d.png", baseDir_.c_str(), layer.name().c_str(), i, j);
|
|
|
|
cv::imwrite(pathBuf, image);
|
|
|
|
data += imagePixels;
|
|
}
|
|
}
|
|
}
|
|
|