Added possibility to load PNG images from a data buffer. (#2137)

* Added possibility to load PNG images from a data buffer.

* Fixed code not compiling with some versions of libpng that doesn't have const specifier.

* Used FileInfo struct as a single parameter for the read_image method.

dlib/image_loader/png_loader.cpp

@@ -29,12 +29,40 @@ namespace dlib
         png_infop end_info_;
     };
 
+    struct PngBufferReaderState
+    {
+        const unsigned char* buffer_;
+        size_t buffer_size_;
+        size_t current_pos_;
+    };
+
+    struct FileInfo
+    {
+        FileInfo( FILE *fp, const char* filename ) : fp_( fp ), filename_( filename )
+        {
+        }
+
+        FileInfo( const unsigned char* buffer, size_t buffer_size ) : buffer_( buffer ), buffer_size_( buffer_size )
+        {
+        }
+
+        ~FileInfo()
+        {
+            if ( fp_ != nullptr ) fclose( fp_ );
+        }
+
+        FILE* fp_{nullptr};
+        const char* filename_{nullptr};
+        const unsigned char* buffer_{nullptr};
+        size_t buffer_size_{0};
+    };
+
 // ----------------------------------------------------------------------------------------
 
     png_loader::
     png_loader( const char* filename ) : height_( 0 ), width_( 0 )
     {
-        read_image( filename );
+        read_image( check_file( filename ) );
     }
 
 // ----------------------------------------------------------------------------------------
@@ -42,7 +70,7 @@ namespace dlib
     png_loader::
     png_loader( const std::string& filename ) : height_( 0 ), width_( 0 )
     {
-        read_image( filename.c_str() );
+        read_image( check_file( filename.c_str() ) );
     }
 
 // ----------------------------------------------------------------------------------------
@@ -50,7 +78,15 @@ namespace dlib
     png_loader::
     png_loader( const dlib::file& f ) : height_( 0 ), width_( 0 )
     {
-        read_image( f.full_name().c_str() );
+        read_image( check_file( f.full_name().c_str() ) );
+    }
+
+// ----------------------------------------------------------------------------------------
+
+    png_loader::
+    png_loader( const unsigned char* image_buffer, size_t buffer_size ) : height_( 0 ), width_( 0 )
+    {
+        read_image( std::unique_ptr<FileInfo>( new FileInfo( image_buffer, buffer_size ) ) );
     }
 
 // ----------------------------------------------------------------------------------------
@@ -96,6 +132,23 @@ namespace dlib
         return ( color_type_ == PNG_COLOR_TYPE_RGB_ALPHA );
     }
 
+// ----------------------------------------------------------------------------------------
+    
+    std::unique_ptr<FileInfo> png_loader::check_file( const char* filename )
+    {
+        if ( filename == NULL )
+        {
+            throw image_load_error("png_loader: invalid filename, it is NULL");
+        }
+        FILE *fp = fopen( filename, "rb" );
+        if ( !fp )
+        {
+            throw image_load_error(std::string("png_loader: unable to open file ") + filename);
+        }
+
+        return std::unique_ptr<FileInfo>( new FileInfo( fp, filename ) );
+    }
+
 // ----------------------------------------------------------------------------------------
 
     // Don't do anything when libpng calls us to tell us about an error.  Just return to 
@@ -108,35 +161,63 @@ namespace dlib
     {
     }
 
-    void png_loader::read_image( const char* filename )
+    void png_buffer_reader(png_structp png_ptr, png_bytep data, size_t length) 
     {
-        ld_.reset(new LibpngData);
-        if ( filename == NULL )
-        {
-            throw image_load_error("png_loader: invalid filename, it is NULL");
-        }
-        FILE *fp = fopen( filename, "rb" );
-        if ( !fp )
+        PngBufferReaderState* state = static_cast<PngBufferReaderState*>( png_get_io_ptr( png_ptr ) );
+        if ( length > ( state->buffer_size_ - state->current_pos_ ) )
         {
-            throw image_load_error(std::string("png_loader: unable to open file ") + filename);
+            png_error(png_ptr, "png_loader: read error in png_buffer_reader");
         }
-        png_byte sig[8];
-        if (fread( sig, 1, 8, fp ) != 8)
+        memcpy( data, state->buffer_ + state->current_pos_, length );
+        state->current_pos_ += length;
+    }
+
+    void png_loader::read_image( std::unique_ptr<FileInfo> file_info )
+    {
+        if ( !file_info ) throw image_load_error(std::string("png_loader: invalid file_info, it is NULL"));
+
+        ld_.reset(new LibpngData);
+
+        constexpr png_size_t png_header_size = 8;
+        std::string load_error_info;
+        
+        if ( file_info->fp_ != NULL )
         {
-            fclose( fp );
-            throw image_load_error(std::string("png_loader: error reading file ") + filename);
+            png_byte sig[png_header_size];
+            if (fread( sig, 1, png_header_size, file_info->fp_ ) != png_header_size)
+            {
+                throw image_load_error(std::string("png_loader: error reading file ") + file_info->filename_);
+            }
+            load_error_info = std::string(" in file ") + file_info->filename_;
+            if ( png_sig_cmp( sig, 0, png_header_size ) != 0 )
+            {
+                throw image_load_error(std::string("png_loader: format error") + load_error_info);
+            }
         }
-        if ( png_sig_cmp( sig, 0, 8 ) != 0 )
+        else
         {
-            fclose( fp );
-            throw image_load_error(std::string("png_loader: format error in file ") + filename);
+            if ( file_info->buffer_ == NULL )
+            {
+                throw image_load_error(std::string("png_loader: invalid image buffer, it is NULL"));
+            }
+            if ( file_info->buffer_size_ == 0 )
+            {
+                throw image_load_error(std::string("png_loader: invalid image buffer size, it is 0"));
+            }
+            if ( file_info->buffer_size_ < png_header_size ||
+                 png_sig_cmp( (png_bytep)file_info->buffer_, 0, png_header_size ) != 0 )
+            {
+                throw image_load_error(std::string("png_loader: format error in image buffer"));
+            }
+
+            buffer_reader_state_.reset(new PngBufferReaderState);
         }
+        
         ld_->png_ptr_ = png_create_read_struct( PNG_LIBPNG_VER_STRING, NULL, &png_loader_user_error_fn_silent, &png_loader_user_warning_fn_silent );
         if ( ld_->png_ptr_ == NULL )
         {
-            fclose( fp );
             std::ostringstream sout;
-            sout << "Error, unable to allocate png structure while opening file " << filename << std::endl;
+            sout << "Error, unable to allocate png structure" << std::endl;
             const char* runtime_version = png_get_header_ver(NULL);
             if (runtime_version && std::strcmp(PNG_LIBPNG_VER_STRING, runtime_version) != 0)
             {
@@ -149,30 +230,39 @@ namespace dlib
         ld_->info_ptr_ = png_create_info_struct( ld_->png_ptr_ );
         if ( ld_->info_ptr_ == NULL )
         {
-            fclose( fp );
             png_destroy_read_struct( &( ld_->png_ptr_ ), ( png_infopp )NULL, ( png_infopp )NULL );
-            throw image_load_error(std::string("png_loader: parse error in file ") + filename);
+            throw image_load_error(std::string("png_loader: unable to allocate png info structure") + load_error_info);
         }
         ld_->end_info_ = png_create_info_struct( ld_->png_ptr_ );
         if ( ld_->end_info_ == NULL )
         {
-            fclose( fp );
             png_destroy_read_struct( &( ld_->png_ptr_ ), &( ld_->info_ptr_ ), ( png_infopp )NULL );
-            throw image_load_error(std::string("png_loader: parse error in file ") + filename);
+            throw image_load_error(std::string("png_loader: unable to allocate png info structure") + load_error_info);
         }
 
         if (setjmp(png_jmpbuf(ld_->png_ptr_)))
         {
             // If we get here, we had a problem writing the file 
-            fclose(fp);
             png_destroy_read_struct( &( ld_->png_ptr_ ), &( ld_->info_ptr_ ), &( ld_->end_info_ ) );
-            throw image_load_error(std::string("png_loader: parse error in file ") + filename);
+            throw image_load_error(std::string("png_loader: parse error") + load_error_info);
         }
 
         png_set_palette_to_rgb(ld_->png_ptr_);
 
-        png_init_io( ld_->png_ptr_, fp );
-        png_set_sig_bytes( ld_->png_ptr_, 8 );
+        if ( file_info->fp_ != NULL )
+        {
+            png_init_io( ld_->png_ptr_, file_info->fp_ );
+        }
+        else
+        {
+            buffer_reader_state_->buffer_ = file_info->buffer_;
+            buffer_reader_state_->buffer_size_ = file_info->buffer_size_;
+            // skipping header
+            buffer_reader_state_->current_pos_ = png_header_size;
+            png_set_read_fn( ld_->png_ptr_, buffer_reader_state_.get(), png_buffer_reader );
+        }
+
+        png_set_sig_bytes( ld_->png_ptr_, png_header_size );
         // flags force one byte per channel output
         byte_orderer bo;
         int png_transforms = PNG_TRANSFORM_PACKING;
@@ -190,25 +280,22 @@ namespace dlib
             color_type_ != PNG_COLOR_TYPE_RGB_ALPHA &&
             color_type_ != PNG_COLOR_TYPE_GRAY_ALPHA)
         {
-            fclose( fp );
             png_destroy_read_struct( &( ld_->png_ptr_ ), &( ld_->info_ptr_ ), &( ld_->end_info_ ) );
-            throw image_load_error(std::string("png_loader: unsupported color type in file ") + filename);
+            throw image_load_error(std::string("png_loader: unsupported color type") + load_error_info);
         }
 
         if (bit_depth_ != 8 && bit_depth_ != 16)
         {
-            fclose( fp );
             png_destroy_read_struct( &( ld_->png_ptr_ ), &( ld_->info_ptr_ ), &( ld_->end_info_ ) );
-            throw image_load_error("png_loader: unsupported bit depth of " + cast_to_string(bit_depth_) + " in file " + std::string(filename));
+            throw image_load_error("png_loader: unsupported bit depth of " + cast_to_string(bit_depth_) + load_error_info);
         }
 
         ld_->row_pointers_ = png_get_rows( ld_->png_ptr_, ld_->info_ptr_ );
 
-        fclose( fp );
         if ( ld_->row_pointers_ == NULL )
         {
             png_destroy_read_struct( &( ld_->png_ptr_ ), &( ld_->info_ptr_ ), &( ld_->end_info_ ) );
-            throw image_load_error(std::string("png_loader: parse error in file ") + filename);
+            throw image_load_error(std::string("png_loader: parse error") + load_error_info);
         }
     }
 

dlib/image_loader/png_loader.h

@@ -15,6 +15,8 @@ namespace dlib
 {
 
     struct LibpngData;
+    struct PngBufferReaderState;
+    struct FileInfo;
     class png_loader : noncopyable
     {
     public:
@@ -22,6 +24,7 @@ namespace dlib
         png_loader( const char* filename );
         png_loader( const std::string& filename );
         png_loader( const dlib::file& f );
+        png_loader( const unsigned char* image_buffer, size_t buffer_size );
         ~png_loader();
 
         bool is_gray() const;
@@ -191,11 +194,13 @@ namespace dlib
 
     private:
         const unsigned char* get_row( unsigned i ) const;
-        void read_image( const char* filename );
+        std::unique_ptr<FileInfo> check_file( const char* filename );
+        void read_image( std::unique_ptr<FileInfo> file_info );
         unsigned height_, width_;
         unsigned bit_depth_;
         int color_type_;
         std::unique_ptr<LibpngData> ld_;
+        std::unique_ptr<PngBufferReaderState> buffer_reader_state_;
     };
 
 // ----------------------------------------------------------------------------------------
@@ -211,6 +216,31 @@ namespace dlib
         png_loader(file_name).get_image(image);
     }
 
+    template <
+        typename image_type
+        >
+    void load_png (
+        image_type& image,
+        const unsigned char* image_buffer,
+        size_t buffer_size
+    )
+    {
+        png_loader(image_buffer, buffer_size).get_image(image);
+    }
+
+    template <
+        typename image_type
+        >
+    void load_png (
+        image_type& image,
+        const char* image_buffer,
+        size_t buffer_size
+    )
+    {
+        png_loader(reinterpret_cast<const unsigned char*>(image_buffer), buffer_size).get_image(image);
+    }
+
+
 // ----------------------------------------------------------------------------------------
 
 }

dlib/image_loader/png_loader_abstract.h

@@ -65,6 +65,19 @@ namespace dlib
                   us from loading the given PNG file.
         !*/
 
+        png_loader( 
+            const unsigned char* image_buffer,
+            size_t buffer_size
+        );
+        /*!
+            ensures
+                - loads the PNG from memory image_buffer of size buffer_size into this object
+            throws
+                - image_load_error
+                  This exception is thrown if there is some error that prevents
+                  us from loading the given PNG buffer.
+        !*/
+
         ~png_loader(
         );
         /*!
@@ -153,6 +166,38 @@ namespace dlib
             - performs: png_loader(file_name).get_image(image);
     !*/
 
+    template <
+        typename image_type
+        >
+    void load_png (
+        image_type& image,
+        const unsigned char* image_buffer,
+        size_t buffer_size
+    );
+    /*!
+        requires
+            - image_type == an image object that implements the interface defined in
+              dlib/image_processing/generic_image.h 
+        ensures
+            - performs: png_loader(image_buffer, buffer_size).get_image(image);
+    !*/
+
+    template <
+        typename image_type
+        >
+    void load_png (
+        image_type& image,
+        const char* image_buffer,
+        size_t buffer_size
+    );
+    /*!
+        requires
+            - image_type == an image object that implements the interface defined in
+              dlib/image_processing/generic_image.h 
+        ensures
+            - performs: png_loader((unsigned char*)image_buffer, buffer_size).get_image(image);
+    !*/
+
 // ----------------------------------------------------------------------------------------
 
 }