/* * 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. }