/*
* save function
*
*/
int libjpeg_(Main_save)(lua_State *L) {
unsigned char *inmem = NULL; /* destination memory (if saving to memory) */
unsigned long inmem_size = 0; /* destination memory size (bytes) */
/* get args */
const char *filename = luaL_checkstring(L, 1);
THTensor *tensor = luaT_checkudata(L, 2, torch_Tensor);
THTensor *tensorc = THTensor_(newContiguous)(tensor);
real *tensor_data = THTensor_(data)(tensorc);
const int save_to_file = luaL_checkint(L, 3);
THByteTensor* tensor_dest = NULL;
if (save_to_file == 0) {
tensor_dest = luaT_checkudata(L, 5, "torch.ByteTensor");
}
int quality = luaL_checkint(L, 4);
if (quality < 0 || quality > 100) {
luaL_error(L, "quality should be between 0 and 100");
}
/* jpeg struct */
struct jpeg_compress_struct cinfo;
struct jpeg_error_mgr jerr;
/* pointer to raw image */
unsigned char *raw_image = NULL;
/* dimensions of the image we want to write */
int width=0, height=0, bytes_per_pixel=0;
int color_space=0;
if (tensorc->nDimension == 3) {
bytes_per_pixel = tensorc->size[0];
height = tensorc->size[1];
width = tensorc->size[2];
if (bytes_per_pixel == 3) {
color_space = JCS_RGB;
} else if (bytes_per_pixel == 1) {
color_space = JCS_GRAYSCALE;
} else {
luaL_error(L, "tensor should have 1 or 3 channels (gray or RGB)");
}
} else if (tensorc->nDimension == 2) {
bytes_per_pixel = 1;
height = tensorc->size[0];
width = tensorc->size[1];
color_space = JCS_GRAYSCALE;
} else {
luaL_error(L, "supports only 1 or 3 dimension tensors");
}
/* alloc raw image data */
raw_image = (unsigned char *)malloc((sizeof (unsigned char))*width*height*bytes_per_pixel);
/* convert tensor to raw bytes */
int x,y,k;
for (k=0; k<bytes_per_pixel; k++) {
for (y=0; y<height; y++) {
for (x=0; x<width; x++) {
raw_image[(y*width+x)*bytes_per_pixel+k] = *tensor_data++;
}
}
}
/* this is a pointer to one row of image data */
JSAMPROW row_pointer[1];
FILE *outfile = NULL;
if (save_to_file == 1) {
outfile = fopen( filename, "wb" );
if ( !outfile ) {
luaL_error(L, "Error opening output jpeg file %s\n!", filename );
}
}
cinfo.err = jpeg_std_error( &jerr );
jpeg_create_compress(&cinfo);
/* specify data source (eg, a file) */
if (save_to_file == 1) {
jpeg_stdio_dest(&cinfo, outfile);
} else {
jpeg_mem_dest(&cinfo, &inmem, &inmem_size);
}
/* Setting the parameters of the output file here */
cinfo.image_width = width;
cinfo.image_height = height;
cinfo.input_components = bytes_per_pixel;
cinfo.in_color_space = color_space;
/* default compression parameters, we shouldn't be worried about these */
jpeg_set_defaults( &cinfo );
jpeg_set_quality(&cinfo, quality, (boolean)0);
/* Now do the compression .. */
jpeg_start_compress( &cinfo, TRUE );
/* like reading a file, this time write one row at a time */
while( cinfo.next_scanline < cinfo.image_height ) {
row_pointer[0] = &raw_image[ cinfo.next_scanline * cinfo.image_width * cinfo.input_components];
jpeg_write_scanlines( &cinfo, row_pointer, 1 );
}
/* similar to read file, clean up after we're done compressing */
jpeg_finish_compress( &cinfo );
jpeg_destroy_compress( &cinfo );
if (outfile != NULL) {
fclose( outfile );
}
if (save_to_file == 0) {
THByteTensor_resize1d(tensor_dest, inmem_size); /* will fail if it's not a Byte Tensor */
unsigned char* tensor_dest_data = THByteTensor_data(tensor_dest);
memcpy(tensor_dest_data, inmem, inmem_size);
free(inmem);
}
/* some cleanup */
free(raw_image);
THTensor_(free)(tensorc);
/* success code is 1! */
return 1;
}
static int torchzfp_(Main_compress)(lua_State *L) {
THTensor* in = reinterpret_cast<THTensor*>(
luaT_checkudata(L, 1, torch_Tensor));
real* in_data = THTensor_(data)(in);
const uint32_t dim = in->nDimension;
if (dim == 0) {
luaL_error(L, "Input tensor must not be empty");
}
THByteTensor* out = reinterpret_cast<THByteTensor*>(
luaT_checkudata(L, 2, "torch.ByteTensor"));
const double accuracy = static_cast<double>(lua_tonumber(L, 3));
// Hacky code to figure out what type 'real' is at runtime. This really should
// be template specialization (so it's compiled in at runtime).
real dummy;
static_cast<void>(dummy); // Silence compiler warnings.
zfp_type type;
if (typeid(dummy) == typeid(float)) {
type = zfp_type_float;
} else if (typeid(dummy) == typeid(double)) {
type = zfp_type_double;
} else {
luaL_error(L, "Input type must be double or float.");
}
// Allocate meta data for the array.
zfp_field* field;
uint32_t dim_zfp;
if (dim == 1) {
field = zfp_field_1d(in_data, type, in->size[0]);
dim_zfp = 1;
} else if (dim == 2) {
field = zfp_field_2d(in_data, type, in->size[1], in->size[0]);
dim_zfp = 2;
} else if (dim == 3) {
field = zfp_field_3d(in_data, type, in->size[2], in->size[1], in->size[0]);
dim_zfp = 3;
} else {
// ZFP only allows up to 3D tensors, so we'll have to treat the input
// tensor as a concatenated 3D tensor. This will affect compression ratios
// but there's not much we can do about this.
uint32_t sizez = 1;
for (uint32_t i = 0; i < dim - 2; i++) {
sizez *= in->size[i];
}
uint32_t sizey = in->size[dim - 2];
uint32_t sizex = in->size[dim - 1];
field = zfp_field_3d(in_data, type, sizex, sizey, sizez);
dim_zfp = 4;
}
// Allocate meta data for the compressed stream.
zfp_stream* zfp = zfp_stream_open(NULL);
// Set stream compression mode and parameters.
zfp_stream_set_accuracy(zfp, accuracy, type);
// Allocate buffer for compressed data.
size_t bufsize = zfp_stream_maximum_size(zfp, field);
std::unique_ptr<uint8_t[]> buffer(new uint8_t[bufsize]);
// Associate bit stream with allocated buffer.
bitstream* stream = stream_open(buffer.get(), bufsize);
zfp_stream_set_bit_stream(zfp, stream);
zfp_stream_rewind(zfp);
// Compress entire array.
const size_t zfpsize = zfp_compress(zfp, field);
// Clean up.
zfp_field_free(field);
zfp_stream_close(zfp);
stream_close(stream);
if (!zfpsize) {
luaL_error(L, "ZFP compression failed!");
}
// Copy the compressed array into the return tensor. NOTE: Torch does not
// support in-place resize with shrink. If you resize smaller you ALWAYS
// keep around the memory, so unfortuantely this copy is necessary (i.e.
// we will always need to perform the compression in a temporary buffer
// first).
THByteTensor_resize1d(out, zfpsize);
unsigned char* out_data = THByteTensor_data(out);
memcpy(out_data, buffer.get(), zfpsize);
return 0; // Recall: number of lua return items.
}
