Homec4science
Diffusion TV norm inpainting cuda phpc2017 (master) Browse

tvnorminpaintcudaphpc2017/dataloader.cpp
93358a8efc04master

Last Change
Raw File

dataloader.cpp
View Options

#include "dataloader.h"
#include <vector>
#ifdef USE_OPEN_EXR
#include <OpenEXR/ImfInputFile.h>
#include <OpenEXR/ImfOutputFile.h>
#include <OpenEXR/ImfIO.h>
#include <OpenEXR/ImfFrameBuffer.h>
#include <OpenEXR/ImathBox.h>
#include <OpenEXR/ImfChannelList.h>
#else
#include <H5Cpp.h>
#endif
using namespace std;
#ifndef USE_OPEN_EXR
using namespace H5;
#endif
#ifdef USE_OPEN_EXR
bool isOpenExrFile (const char fileName[]) {
std::ifstream f (fileName, std::ios_base::binary);
char b[4];
f.read (b, sizeof (b));
return !!f && b[0] == 0x76 && b[1] == 0x2f && b[2] == 0x31 && b[3] == 0x01;
}
#endif
bool readArray2D (const char fileName[],
float* & pixels,
int &width,
int &height,
std::vector<std::string> names) {
#ifdef USE_OPEN_EXR
if (!isOpenExrFile(fileName)) {
return false;
}
Imf::InputFile file (fileName);
std::string cName;
const Imf::ChannelList &channels = file.header().channels();
for (std::string const& cref : names) {
const Imf::Channel *channelPtr = channels.findChannel(cref);
if (channelPtr != nullptr) {
cName = cref;
goto prog_suite;
}
}
return false;
prog_suite:
Imath::Box2i dw = file.header().dataWindow();
width = dw.max.x - dw.min.x + 1;
height = dw.max.y - dw.min.y + 1;
pixels = new float[height*width];
Imf::FrameBuffer frameBuffer;
frameBuffer.insert (cName, // name
Imf::Slice (Imf::FLOAT, // type
(char *) (pixels - // base
dw.min.x -
dw.min.y * width),
sizeof (*pixels) * 1, // xStride
sizeof (*pixels) * width,// yStride
1, 1, // x/y sampling
FLT_MAX)); // fillValue
file.setFrameBuffer (frameBuffer);
file.readPixels (dw.min.y, dw.max.y);
#else
H5File file(fileName, H5F_ACC_RDONLY);
DataSet dataset = file.openDataSet( names[0] );
DataType datatype = dataset.getDataType();
if (!(datatype == H5::PredType::NATIVE_FLOAT)) {
return false;
}
DataSpace dataspace = dataset.getSpace();
int ndims = dataspace.getSimpleExtentNdims();
if (ndims != 2) {
return false;
}
hsize_t dims_out[2];
dataspace.getSimpleExtentDims( dims_out, nullptr);
width = dims_out[0];
height = dims_out[1];
DataSpace memspace(2, dims_out);
memspace.selectAll();
pixels = new float[height*width];
dataset.read( pixels, datatype, memspace, dataspace );
#endif
return true;
}
bool writeFile (const char fileName[],
std::vector<const float*> const& pixels,
std::vector<const char *> names,
int sx,
int sy) {
if (pixels.size() != names.size()) {
return false;
}
#ifdef USE_OPEN_EXR
Imf::Header header (sx, sy);
for (const char* name : names) {
header.channels().insert (name, Imf::Channel (Imf::FLOAT));
}
Imf::OutputFile oFile(fileName, header);
Imf::FrameBuffer frameBuffer;
for (int i = 0; i < names.size(); i++) {
frameBuffer.insert (names[i], Imf::Slice(Imf::FLOAT,
(char *) pixels[i],
sizeof (*(pixels[i])) * 1,
sizeof (*(pixels[i])) * sx));
}
oFile.setFrameBuffer (frameBuffer);
oFile.writePixels (sy);
#else
H5File outfile(fileName, H5F_ACC_TRUNC);
hsize_t dimsf[2];
dimsf[0] = sx;
dimsf[1] = sy;
DataSpace dataspace(2, dimsf);
DataType datatype(H5::PredType::NATIVE_FLOAT);
for (int i = 0; i < names.size(); i++) {
DataSet dataset = outfile.createDataSet(names[i], datatype, dataspace);
dataset.write(pixels[i], datatype);
}
#endif
return true;
}
DataLoader::DataLoader():
_status(true) ,
_redNeedDestroyed(false) ,
_greenNeedDestroyed(false) ,
_blueNeedDestroyed(false) ,
_maskNeedDestroyed(false) {
}
DataLoader::~DataLoader() {
if (_redNeedDestroyed) {
delete [] _redChannel;
}
if (_greenNeedDestroyed) {
delete [] _greenChannel;
}
if (_blueNeedDestroyed) {
delete [] _blueChannel;
}
if (_maskNeedDestroyed) {
delete [] _mask;
}
}
void DataLoader::loadRGB(string const& fileName) {
bool ok;
ok = readArray2D(fileName.c_str(), _redChannel, _imW, _imH, {"R", "RGB.R", "RGBA.R", "Composite.R"});
if (!ok) {
_status = false;
return;
}
_redNeedDestroyed = true;
ok = readArray2D(fileName.c_str(), _greenChannel, _imW, _imH, {"G", "RGB.G", "RGBA.G", "Composite.G"});
if (!ok) {
_status = false;
return;
}
_greenNeedDestroyed = true;
ok = readArray2D(fileName.c_str(), _blueChannel, _imW, _imH, {"B", "RGB.B", "RGBA.B", "Composite.B"});
if (!ok) {
_status = false;
}
_blueNeedDestroyed = true;
}
void DataLoader::loadMask(string const& fileName) {
bool ok;
int w;
int h;
ok = readArray2D(fileName.c_str(), _mask, w, h, {"A", "RGBA.A", "Composite.A"});
if (!ok || w != _imW || h != _imH) {
_status = false;
}
if (ok) {
_maskNeedDestroyed = true;
}
}
void DataLoader::outputRGB(string const& fileName) {
writeFile (fileName.c_str(),
{_redChannel, _greenChannel, _blueChannel},
{"R", "G", "B"},
_imW,
_imH);
}
bool DataLoader::status() const{
return _status;
}
float* DataLoader::redChannel() {
return _redChannel;
}
float* DataLoader::greenChannel() {
return _greenChannel;
}
float* DataLoader::blueChannel() {
return _blueChannel;
}
float* DataLoader::mask() {
return _mask;
}
int DataLoader::imW() const {
return _imW;
}
int DataLoader::imH() const {
return _imH;
}
c4science · Help