support abbreviated JPEG data streams

src/decode.jl

@@ -57,11 +57,12 @@ filename = testimage("earth", download_only=true)
 """
 function jpeg_decode(
         ::Type{CT},
+        table::Union{Nothing,Vector{UInt8}},
         data::Vector{UInt8};
         transpose=false,
         scale_ratio::Union{Nothing,Real}=nothing,
         preferred_size::Union{Nothing,Tuple}=nothing) where CT<:Colorant
-    _jpeg_check_bytes(data)
+    _jpeg_check_bytes(table, data)
     out_CT, jpeg_cls = _jpeg_out_color_space(CT)
 
     cinfo_ref = Ref(LibJpeg.jpeg_decompress_struct())
@@ -70,8 +71,19 @@ function jpeg_decode(
         cinfo = cinfo_ref[]
         cinfo.err = LibJpeg.jpeg_std_error(jerr)
         LibJpeg.jpeg_create_decompress(cinfo_ref)
+        if !(isnothing(table) || isempty(table))
+            # if abbreviated table is provided, read it first
+            LibJpeg.jpeg_mem_src(cinfo_ref, table, length(table))
+            v = LibJpeg.jpeg_read_header(cinfo_ref, false)
+            if v != LibJpeg.JPEG_HEADER_TABLES_ONLY
+                @warn "expected a header-only JPEG data stream" v
+            end
+        end
         LibJpeg.jpeg_mem_src(cinfo_ref, data, length(data))
-        LibJpeg.jpeg_read_header(cinfo_ref, true)
+        v = LibJpeg.jpeg_read_header(cinfo_ref, true)
+        if v != LibJpeg.JPEG_HEADER_OK
+            error("expected a valid JPEG data stream")
+        end
 
         # set decompression parameters, if given
         if !isnothing(preferred_size) && !transpose
@@ -106,19 +118,31 @@ function jpeg_decode(
         LibJpeg.jpeg_destroy_decompress(cinfo_ref)
     end
 end
-jpeg_decode(data; kwargs...) = jpeg_decode(_default_out_color_space(data), data; kwargs...)
+jpeg_decode(::Type{CT}, data::Vector{UInt8}; kwargs...) where CT = jpeg_decode(CT, nothing, data; kwargs...)
+
+jpeg_decode(table, data; kwargs...) = jpeg_decode(_default_out_color_space(table, data), data; kwargs...)
+jpeg_decode(data; kwargs...) = jpeg_decode(nothing, data; kwargs...)
 
 # TODO(johnnychen94): support Progressive JPEG
 # TODO(johnnychen94): support partial decoding
+function jpeg_decode(::Type{CT}, table_file::AbstractString, data_file::AbstractString; kwargs...) where CT<:Colorant
+    jpeg_decode(CT, read(table_file), read(data_file); kwargs...)
+end
 function jpeg_decode(::Type{CT}, filename::AbstractString; kwargs...) where CT<:Colorant
-    open(filename, "r") do io
-        jpeg_decode(CT, io; kwargs...)
-    end
+    jpeg_decode(CT, nothing, read(filename); kwargs...)
+end
+function jpeg_decode(table_file::AbstractString, data_file::AbstractString; kwargs...)
+    jpeg_decode(read(table_file), read(filename); kwargs...)
 end
-jpeg_decode(filename::AbstractString; kwargs...) = jpeg_decode(read(filename); kwargs...)
+jpeg_decode(filename::AbstractString; kwargs...) = jpeg_decode(nothing, read(filename); kwargs...)
+
+jpeg_decode(table_io::IO, data_io::IO; kwargs...) = jpeg_decode(read(table_io), read(data_io); kwargs...)
+jpeg_decode(io::IO; kwargs...) = jpeg_decode(nothing, read(io); kwargs...)
 
-jpeg_decode(io::IO; kwargs...) = jpeg_decode(read(io); kwargs...)
-jpeg_decode(::Type{CT}, io::IO; kwargs...) where CT<:Colorant = jpeg_decode(CT, read(io); kwargs...)
+function jpeg_decode(::Type{CT}, table_io::IO, data_io::IO; kwargs...) where CT<:Colorant
+    jpeg_decode(CT, read(table_io), read(data_io); kwargs...)
+end
+jpeg_decode(::Type{CT}, io::IO; kwargs...) where CT<:Colorant = jpeg_decode(CT, nothing, read(io); kwargs...)
 
 function _jpeg_decode!(out::Matrix{<:Colorant}, cinfo_ref::Ref{LibJpeg.jpeg_decompress_struct})
     row_stride = size(out, 1) * length(eltype(out))
@@ -175,13 +199,18 @@ function _cal_scale_ratio(::Nothing, preferred_size::Tuple, cinfo)
     return _allowed_scale_ratios[idx]
 end
 
-function _default_out_color_space(data::Vector{UInt8})
-    _jpeg_check_bytes(data)
+function _default_out_color_space(table, data::Vector{UInt8})
+    _jpeg_check_bytes(table, data)
     cinfo_ref = Ref(LibJpeg.jpeg_decompress_struct())
     try
         jerr = Ref{LibJpeg.jpeg_error_mgr}()
         cinfo_ref[].err = LibJpeg.jpeg_std_error(jerr)
         LibJpeg.jpeg_create_decompress(cinfo_ref)
+        if !(isnothing(table) || isempty(table))
+            # if abbreviated table is provided, read it first
+            LibJpeg.jpeg_mem_src(cinfo_ref, table, length(table))
+            LibJpeg.jpeg_read_header(cinfo_ref, false)
+        end
         LibJpeg.jpeg_mem_src(cinfo_ref, data, length(data))
         LibJpeg.jpeg_read_header(cinfo_ref, true)
         LibJpeg.jpeg_calc_output_dimensions(cinfo_ref)
@@ -202,8 +231,8 @@ end
 
 # provides some basic integrity check
 # TODO(johnnychen94): redirect libjpeg-turbo error to julia
-_jpeg_check_bytes(filename::AbstractString) = open(_jpeg_check_bytes, filename, "r")
-function _jpeg_check_bytes(io::IO)
+_jpeg_check_bytes(filename::AbstractString; kwargs...) = open(_jpeg_check_bytes, filename, "r")
+function _jpeg_check_bytes(io::IO; kwargs...)
     seekend(io)
     nbytes = position(io)
     nbytes > 623 || throw(ArgumentError("Invalid number of bytes."))
@@ -215,9 +244,14 @@ function _jpeg_check_bytes(io::IO)
     append!(buf, read(io, 2))
     return _jpeg_check_bytes(buf)
 end
-function _jpeg_check_bytes(data::Vector{UInt8})
-    length(data) > 623 || throw(ArgumentError("Invalid number of bytes."))
+function _jpeg_check_bytes(data::Vector{UInt8}; check_length=true)
+    check_length && (length(data) > 623 || throw(ArgumentError("Invalid number of bytes.")))
     data[1:2] == [0xff, 0xd8] || throw(ArgumentError("Invalid JPEG byte sequence."))
     data[end-1:end] == [0xff, 0xd9] || @warn "Premature end of JPEG byte sequence."
     return true
 end
+_jpeg_check_bytes(::Nothing; kwargs...) = true
+function _jpeg_check_bytes(table, data)
+    _jpeg_check_bytes(table; check_length=false)
+    _jpeg_check_bytes(data; check_length=true)
+end

src/encode.jl

@@ -49,20 +49,36 @@ function jpeg_encode(img::AbstractMatrix{T}; transpose=false, kwargs...) where T
     # contiguous memeory layout already makes a transpose.
     img = transpose ? convert(AT, img) : convert(AT, PermutedDimsArray(img, (2, 1)))
 
-    return _encode(img; kwargs...)
+    table, data = _encode(img; kwargs...)
+    return isempty(table) ? data : (table, data) # for backward compatibility
 end
 
 function jpeg_encode(filename::AbstractString, img; kwargs...)
     open(filename, "w") do io
         jpeg_encode(io, img; kwargs...)
     end
 end
+function jpeg_encode(table_file::AbstractString, data_file::AbstractString, img; kwargs...)
+    open(table_file, "w") do table_io
+        open(data_file, "w") do data_io
+            jpeg_encode(table_io, data_io, img; kwargs...)
+        end
+    end
+end
 # TODO(johnnychen94): further improve the performance via asynchronously IO and buffer reuse.
 jpeg_encode(io::IO, img; kwargs...) = write(io, jpeg_encode(img; kwargs...))
-
+function jpeg_encode(table_io::IO, data_io::IO, img; kwargs...)
+    table, data = jpeg_encode(img; separate_tables=true, kwargs...)
+    n = write(table_io, table)
+    n += write(data_io, data)
+    return n
+end
 
 function _encode(
     img::Matrix{<:Colorant};
+    separate_tables::Bool = false,
+    skip_tables::Bool = false,
+    skip_data::Bool = false,
     colorspace::Union{Nothing,Type} = nothing,
     # ImageMagick: "the default is to use the estimated quality of your input image if it can be determined, otherwise 92."
     quality::Union{Nothing,Int} = 92,
@@ -76,54 +92,92 @@ function _encode(
     Y_density::Union{Nothing,Int} = nothing,
     write_Adobe_marker::Union{Nothing,Bool} = nothing
 )
-    cinfo = LibJpeg.jpeg_compress_struct()
-    cinfo_ref = Ref(cinfo)
-    jerr = Ref{LibJpeg.jpeg_error_mgr}()
-    cinfo.err = LibJpeg.jpeg_std_error(jerr)
-    LibJpeg.jpeg_create_compress(cinfo_ref)
-
-    # set input image information
-    cinfo.image_width = size(img, 1)
-    cinfo.image_height = size(img, 2)
-    cinfo.input_components = jpeg_components(img)
-    cinfo.in_color_space = jpeg_color_space(img)
-
-    # set compression keywords
-    # it's recommended to call `jpeg_set_defaults` first before setting custom parameters
-    # as it's more likely to provide a working parameters and is more likely to be working
-    # correctly in the future.
-    LibJpeg.jpeg_set_defaults(cinfo_ref)
-    isnothing(colorspace) || LibJpeg.jpeg_set_colorspace(cinfo_ref, jpeg_color_space(colorspace))
-    isnothing(quality) || LibJpeg.jpeg_set_quality(cinfo_ref, quality, true)
-    isnothing(arith_code) || (cinfo.arith_code = arith_code)
-    isnothing(optimize_coding) || (cinfo.optimize_coding = optimize_coding)
-    isnothing(smoothing_factor) || (cinfo.smoothing_factor = smoothing_factor)
-    isnothing(write_JFIF_header) || (cinfo.write_JFIF_header = write_JFIF_header)
-    if !isnothing(JFIF_version)
-        cinfo.JFIF_major_version = UInt8(JFIF_version.major)
-        cinfo.JFIF_minor_version = UInt8(JFIF_version.minor)
-    end
-    isnothing(density_unit) || (cinfo.density_unit = density_unit)
-    isnothing(X_density) || (cinfo.X_density = X_density)
-    isnothing(Y_density) || (cinfo.Y_density = Y_density)
-    isnothing(write_Adobe_marker) || (cinfo.write_Adobe_marker = write_Adobe_marker)
-
-    # set destination
-    # TODO(johnnychen94): allow pre-allocated buffer
-    bufsize = Ref{Culong}(0)
-    buf_ptr = Ref{Ptr{UInt8}}(C_NULL)
-    LibJpeg.jpeg_mem_dest(cinfo_ref, buf_ptr, bufsize)
-
-    # compression stage
-    LibJpeg.jpeg_start_compress(cinfo_ref, true)
-    row_stride = size(img, 1) * jpeg_components(img)
-    row_pointer = Ref{Ptr{UInt8}}(0)
-    while (cinfo.next_scanline < cinfo.image_height)
-        row_pointer[] = pointer(img) + cinfo.next_scanline * row_stride
-        LibJpeg.jpeg_write_scanlines(cinfo_ref, row_pointer, 1);
+    cinfo_ref = Ref(LibJpeg.jpeg_compress_struct())
+    cinfo = cinfo_ref[]
+    try
+        jerr = Ref{LibJpeg.jpeg_error_mgr}()
+        cinfo.err = LibJpeg.jpeg_std_error(jerr)
+        LibJpeg.jpeg_create_compress(cinfo_ref)
+
+        # set input image information
+        cinfo.image_width = size(img, 1)
+        cinfo.image_height = size(img, 2)
+        cinfo.input_components = jpeg_components(img)
+        cinfo.in_color_space = jpeg_color_space(img)
+
+        # set compression keywords
+        # it's recommended to call `jpeg_set_defaults` first before setting custom parameters
+        # as it's more likely to provide a working parameters and is more likely to be working
+        # correctly in the future.
+        LibJpeg.jpeg_set_defaults(cinfo_ref)
+        isnothing(colorspace) || LibJpeg.jpeg_set_colorspace(cinfo_ref, jpeg_color_space(colorspace))
+        isnothing(quality) || LibJpeg.jpeg_set_quality(cinfo_ref, quality, true)
+        isnothing(arith_code) || (cinfo.arith_code = arith_code)
+        isnothing(optimize_coding) || (cinfo.optimize_coding = optimize_coding)
+        isnothing(smoothing_factor) || (cinfo.smoothing_factor = smoothing_factor)
+        isnothing(write_JFIF_header) || (cinfo.write_JFIF_header = write_JFIF_header)
+        if !isnothing(JFIF_version)
+            cinfo.JFIF_major_version = UInt8(JFIF_version.major)
+            cinfo.JFIF_minor_version = UInt8(JFIF_version.minor)
+        end
+        isnothing(density_unit) || (cinfo.density_unit = density_unit)
+        isnothing(X_density) || (cinfo.X_density = X_density)
+        isnothing(Y_density) || (cinfo.Y_density = Y_density)
+        isnothing(write_Adobe_marker) || (cinfo.write_Adobe_marker = write_Adobe_marker)
+
+        if separate_tables
+            # ref: "Abbreviated datastreams and multiple images" section in libjpeg-turbo API Documentation
+            skip_tables && skip_data && error("`separate_tables` is set, but `skip_tables` or `skip_data` is not set.")
+            if skip_tables
+                table_buf = UInt8[]
+            else
+                bufsize = Ref{Culong}(0)
+                table_buf_ptr = Ref{Ptr{UInt8}}(C_NULL)
+                LibJpeg.jpeg_mem_dest(cinfo_ref, table_buf_ptr, bufsize)
+                LibJpeg.jpeg_write_tables(cinfo_ref)
+                table_buf = unsafe_wrap(Array, table_buf_ptr[], bufsize[]; own=true)
+            end
+
+            if skip_data
+                data_buf = UInt8[]
+            else
+                bufsize = Ref{Culong}(0)
+                data_buf_ptr = Ref{Ptr{UInt8}}(C_NULL)
+                LibJpeg.jpeg_mem_dest(cinfo_ref, data_buf_ptr, bufsize)
+                LibJpeg.jpeg_start_compress(cinfo_ref, true)
+                row_stride = size(img, 1) * jpeg_components(img)
+                row_pointer = Ref{Ptr{UInt8}}(0)
+                while (cinfo.next_scanline < cinfo.image_height)
+                    row_pointer[] = pointer(img) + cinfo.next_scanline * row_stride
+                    LibJpeg.jpeg_write_scanlines(cinfo_ref, row_pointer, 1);
+                end
+                LibJpeg.jpeg_finish_compress(cinfo_ref)
+                data_buf = unsafe_wrap(Array, data_buf_ptr[], bufsize[]; own=true)
+            end
+
+            return table_buf, data_buf
+        else
+            if skip_tables || skip_data
+                @warn "`skip_tables` and `skip_data` are not used when `separate_tables` is not set."
+            end
+            # set destination
+            # TODO(johnnychen94): allow pre-allocated buffer
+            bufsize = Ref{Culong}(0)
+            buf_ptr = Ref{Ptr{UInt8}}(C_NULL)
+            LibJpeg.jpeg_mem_dest(cinfo_ref, buf_ptr, bufsize)
+
+            # compression stage
+            LibJpeg.jpeg_start_compress(cinfo_ref, true)
+            row_stride = size(img, 1) * jpeg_components(img)
+            row_pointer = Ref{Ptr{UInt8}}(0)
+            while (cinfo.next_scanline < cinfo.image_height)
+                row_pointer[] = pointer(img) + cinfo.next_scanline * row_stride
+                LibJpeg.jpeg_write_scanlines(cinfo_ref, row_pointer, 1);
+            end
+            LibJpeg.jpeg_finish_compress(cinfo_ref)
+            return UInt8[], unsafe_wrap(Array, buf_ptr[], bufsize[]; own=true)
+        end
+    finally
+        LibJpeg.jpeg_destroy_compress(cinfo_ref)
     end
-    LibJpeg.jpeg_finish_compress(cinfo_ref)
-    LibJpeg.jpeg_destroy_compress(cinfo_ref)
-
-    return unsafe_wrap(Array, buf_ptr[], bufsize[]; own=true)
 end