in tensorflow_io/core/kernels/image_pnm_kernels.cc [25:153]
void Compute(OpKernelContext* context) override {
const Tensor* input_tensor;
OP_REQUIRES_OK(context, context->input("input", &input_tensor));
string input = input_tensor->scalar<tstring>()();
size_t pos = 0;
size_t off = input.find_first_of(" \t\r\n", pos);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no magic"));
string magic = input.substr(pos, off - pos);
OP_REQUIRES(
context,
(magic == "P2" || magic == "P3" || magic == "P5" || magic == "P6"),
errors::InvalidArgument("invalid format: ", magic));
const int64 channels = (magic == "P2" || magic == "P5") ? 1 : 3;
off = input.find_first_not_of(" \t\r\n", off);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no width"));
if (input[off] == '#') {
// comment
while (off < input.size() && input[off] != '\n') {
off++;
}
}
// width, height, max
pos = off;
off = input.find_first_of(" \t\r\n", pos);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no width"));
int64 width;
OP_REQUIRES(context,
(strings::safe_strto64(input.substr(pos, off - pos), &width)),
errors::InvalidArgument("unable to parse width: ",
input.substr(pos, off - pos)));
off = input.find_first_not_of(" \t\r\n", off);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no height"));
pos = off;
off = input.find_first_of(" \t\r\n", pos);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no height"));
int64 height;
OP_REQUIRES(context,
(strings::safe_strto64(input.substr(pos, off - pos), &height)),
errors::InvalidArgument("unable to parse height: ",
input.substr(pos, off - pos)));
off = input.find_first_not_of(" \t\r\n", off);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no max"));
pos = off;
off = input.find_first_of(" \t\r\n", pos);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no max"));
int64 max;
OP_REQUIRES(context,
(strings::safe_strto64(input.substr(pos, off - pos), &max)),
errors::InvalidArgument("unable to parse max: ",
input.substr(pos, off - pos)));
OP_REQUIRES(context, (max == 255 || max == 65535),
errors::InvalidArgument("invalid max value: ", max));
Tensor* image_tensor = nullptr;
OP_REQUIRES_OK(
context, context->allocate_output(
0, TensorShape({height, width, channels}), &image_tensor));
if (magic == "P2" || magic == "P3") {
for (int64 i = 0; i < image_tensor->NumElements(); i++) {
off = input.find_first_not_of(" \t\r\n", off);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("not enough value"));
pos = off;
off = input.find_first_of(" \t\r\n", pos);
OP_REQUIRES(context, (off != string::npos),
errors::InvalidArgument("no value"));
int32 value;
OP_REQUIRES(
context,
(strings::safe_strto32(input.substr(pos, off - pos), &value)),
errors::InvalidArgument("unable to parse value: ",
input.substr(pos, off - pos)));
if (image_tensor->dtype() == DT_UINT8) {
if (max == 255) {
image_tensor->flat<uint8>()(i) = static_cast<uint8>(value);
} else {
image_tensor->flat<uint8>()(i) = static_cast<uint8>(value / 256);
}
} else {
if (max == 255) {
image_tensor->flat<uint16>()(i) = static_cast<uint16>(value * 256);
} else {
image_tensor->flat<uint16>()(i) = static_cast<uint16>(value);
}
}
}
} else {
off++;
if (image_tensor->dtype() == DT_UINT8) {
if (max == 255) {
OP_REQUIRES(context,
(off + image_tensor->NumElements() <= input.size()),
errors::InvalidArgument("not enough data"));
memcpy(image_tensor->flat<uint8>().data(), &input[off],
image_tensor->NumElements());
} else {
// TODO: add support for max = 65535 and dtype = uint8; need test file
OP_REQUIRES(
context, false,
errors::InvalidArgument(
"not supported with max == 65535 and dtype == uint8"));
}
} else {
if (max == 255) {
// TODO: add support for max = 255 and dtype = uint16; need test file
OP_REQUIRES(context, false,
errors::InvalidArgument(
"not supported with max == 255 and dtype == uint16"));
} else {
// network order so switch
for (int64 i = 0; i < image_tensor->NumElements(); i++) {
image_tensor->flat<uint16>()(i) =
static_cast<uint16>((((int32)input[off + i * 2] & 0xFF) << 8) |
(((int32)input[off + i * 2 + 1] & 0xFF)));
}
}
}
}
}