feat: initial jpeg_decode implementation

Author: johnnychen94

docs/src/reference.md

@@ -6,6 +6,7 @@ CurrentModule = JpegTurbo
 
 ```@docs
 jpeg_encode
+jpeg_decode
 ```
 
 ## Utilities

src/JpegTurbo.jl

@@ -8,7 +8,8 @@ include("libjpeg_utils.jl")
 
 include("common.jl")
 include("encode.jl")
+include("decode.jl")
 
-export jpeg_encode
+export jpeg_encode, jpeg_decode
 
 end

src/common.jl

@@ -10,3 +10,25 @@ jpeg_color_space(::Type{CT}) where CT<:RGBA = LibJpeg.JCS_EXT_RGBA
 jpeg_color_space(::Type{CT}) where CT<:BGRA = LibJpeg.JCS_EXT_BGRA
 jpeg_color_space(::Type{CT}) where CT<:ABGR = LibJpeg.JCS_EXT_ABGR
 jpeg_color_space(::Type{CT}) where CT<:ARGB = LibJpeg.JCS_EXT_ARGB
+
+function jpeg_color_space(v::LibJpeg.J_COLOR_SPACE)::Type{<:Colorant}
+    if v === LibJpeg.JCS_GRAYSCALE
+        return Gray{N0f8}
+    elseif v === LibJpeg.JCS_RGB
+        return RGB{N0f8}
+    elseif v === LibJpeg.JCS_EXT_RGB
+        return RGB{N0f8}
+    elseif v === LibJpeg.JCS_EXT_BGR
+        return BGR{N0f8}
+    elseif v === LibJpeg.JCS_EXT_RGBA
+        return RGBA{N0f8}
+    elseif v === LibJpeg.JCS_EXT_BGRA
+        return BGRA{N0f8}
+    elseif v === LibJpeg.JCS_EXT_ABGR
+        return ABGR{N0f8}
+    elseif v === LibJpeg.JCS_EXT_ARGB
+        return ARGB{N0f8}
+    else
+        throw(ArgumentError("Unsupported colorspace $v"))
+    end
+end

src/decode.jl

@@ -0,0 +1,147 @@
+"""
+    jpeg_decode([T,] filename::AbstractString; kwargs...) -> Matrix{T}
+
+Decode the JPEG image from given I/O stream as colorant matrix.
+
+# parameters
+
+- `transpose::Bool`: whether we need to permute the image's width and height dimension
+  before encoding. The default value is `false`.
+- `scale_ratio::Real`: scale the image by ratio `scale_ratio` in `M/8` with `M ∈ 1:16`. The
+  default value is `1`. For values are not in the range, they will be mapped to the nearest
+  value, e.g., `0.3 => 2/8` and `0.35 => 3/8`.
+
+# Examples
+
+```jldoctest
+julia> using JpegTurbo, TestImages, ImageCore
+
+julia> filename = testimage("earth", download_only=true);
+
+julia> img = jpeg_decode(filename); summary(img)
+"3002×3000 Array{RGB{N0f8},2} with eltype RGB{N0f8}"
+
+julia> img = jpeg_decode(Gray, filename; scale_ratio=0.25); summary(img)
+"751×750 Array{Gray{N0f8},2} with eltype Gray{N0f8}"
+```
+
+For image preview and similar purposes, `T` and `scale_ratio` are useful parameters to
+accelerate the JPEG decoding process. For color JPEG image, `jpeg_decode(Gray, filename)` is
+faster than `jpeg_decode(filename)` since the color components need not be processed.
+Smaller `scale_ratio` permits significantly faster decoding since fewer pixels need be
+processed and a simpler IDCT method can be used.
+
+```julia
+using BenchmarkTools, TestImages, JpegTurbo
+filename = testimage("earth", download_only=true)
+# full decompression
+@btime jpeg_decode(filename); # 224.760 ms (7 allocations: 51.54 MiB)
+# only decompress luminance component
+@btime jpeg_decode(Gray, filename); # 91.157 ms (6 allocations: 17.18 MiB)
+# process only a few pixels
+@btime jpeg_decode(filename; scale_ratio=0.25); # 77.254 ms (8 allocations: 3.23 MiB)
+# process only a few pixels for luminance component
+@btime jpeg_decode(Gray, filename; scale_ratio=0.25); # 63.119 ms (6 allocations: 1.08 MiB)
+```
+
+"""
+function jpeg_decode(
+        ::Type{CT},
+        filename::AbstractString;
+        transpose=false,
+        scale_ratio=1) where CT<:Colorant
+    infile = ccall(:fopen, Libc.FILE, (Cstring, Cstring), filename, "rb")
+    @assert infile.ptr != C_NULL
+    out_CT, jpeg_cls = _jpeg_out_color_space(CT)
+
+    local cinfo, out
+    try
+        cinfo = LibJpeg.jpeg_decompress_struct()
+        cinfo_ref = Ref(cinfo)
+        jerr = Ref{LibJpeg.jpeg_error_mgr}()
+        cinfo.err = LibJpeg.jpeg_std_error(jerr)
+        LibJpeg.jpeg_create_decompress(cinfo_ref)
+        LibJpeg.jpeg_stdio_src(cinfo_ref, infile)
+        LibJpeg.jpeg_read_header(cinfo_ref, true)
+
+        # set decompression parameters, if given
+        r = _cal_scale_ratio(scale_ratio)
+        cinfo.scale_num, cinfo.scale_denom = r.num, r.den
+        cinfo.out_color_space = jpeg_cls
+
+        # eagerly calculate dimension information so that `output_XXX` fields are valid.
+        LibJpeg.jpeg_calc_output_dimensions(cinfo_ref)
+        out_size = (Int(cinfo.output_width), Int(cinfo.output_height))
+        if !all(x -> x<=65535, out_size)
+            error("Suspicious inferred image size $out_size: each dimension is expected to have at most 65535 pixels.")
+        end
+        out_ndims = Int(cinfo.output_components)
+        @assert out_ndims == length(out_CT) "Suspicous output color space: $cinfo.out_color_space"
+
+        out = Matrix{out_CT}(undef, out_size)
+        _jpeg_decode!(out, cinfo)
+    finally
+        LibJpeg.jpeg_destroy_decompress(Ref(cinfo))
+        ccall(:fclose, Cint, (Ptr{Libc.FILE},), infile)
+    end
+
+    if out_CT <: CT
+        return transpose ? out : permutedims(out, (2, 1))
+    else
+        return transpose ? CT.(out) : CT.(PermutedDimsArray(out, (2, 1)))
+    end
+end
+function jpeg_decode(filename::AbstractString; kwargs...)
+    return jpeg_decode(_default_out_color_space(filename), filename; kwargs...)
+end
+
+function _jpeg_decode!(out::Matrix{<:Colorant}, cinfo::LibJpeg.jpeg_decompress_struct)
+    cinfo_ref = Ref(cinfo)
+
+    row_stride = size(out, 1) * length(eltype(out))
+    buf = Vector{UInt8}(undef, row_stride)
+    out_uint8 = reinterpret(UInt8, out)
+
+    LibJpeg.jpeg_start_decompress(cinfo_ref)
+    while cinfo.output_scanline < cinfo.output_height
+        # TODO(johnnychen94): try if we can directly write to `out` without using `buf`
+        GC.@preserve buf LibJpeg.jpeg_read_scanlines(cinfo_ref, Ref(pointer(buf)), 1)
+        copyto!(out_uint8, (cinfo.output_scanline-1) * row_stride + 1, buf, 1, row_stride)
+    end
+    LibJpeg.jpeg_finish_decompress(cinfo_ref)
+
+    return out
+end
+
+# libjpeg-turbo only supports ratio M/8 with M from 1 to 16
+const _allowed_scale_ratios = ntuple(i->i//8, 16)
+_cal_scale_ratio(r::Real) = _allowed_scale_ratios[findmin(x->abs(x-r), _allowed_scale_ratios)[2]]
+
+function _default_out_color_space(filename::AbstractString)
+    infile = ccall(:fopen, Libc.FILE, (Cstring, Cstring), filename, "rb")
+    @assert infile.ptr != C_NULL
+    local cinfo
+    try
+        cinfo = LibJpeg.jpeg_decompress_struct()
+        cinfo_ref = Ref(cinfo)
+        jerr = Ref{LibJpeg.jpeg_error_mgr}()
+        cinfo.err = LibJpeg.jpeg_std_error(jerr)
+        LibJpeg.jpeg_create_decompress(cinfo_ref)
+        LibJpeg.jpeg_stdio_src(cinfo_ref, infile)
+        LibJpeg.jpeg_read_header(cinfo_ref, true)
+        LibJpeg.jpeg_calc_output_dimensions(cinfo_ref)
+        return jpeg_color_space(cinfo.out_color_space)
+    finally
+        LibJpeg.jpeg_destroy_decompress(Ref(cinfo))
+        ccall(:fclose, Cint, (Ptr{Libc.FILE},), infile)
+    end
+end
+
+function _jpeg_out_color_space(::Type{CT}) where CT
+    try
+        n0f8(CT), jpeg_color_space(n0f8(CT))
+    catch e
+        @debug "Unsupported libjpeg-turbo color space, fallback to RGB{N0f8}" e
+        RGB{N0f8}, jpeg_color_space(RGB{N0f8})
+    end
+end

test/runtests.jl

@@ -13,8 +13,7 @@ function decode_encode(img; kwargs...)
     tmpfile = joinpath(tmpdir, "tmp.jpg")
     buf = @inferred jpeg_encode(img; kwargs...)
     write(tmpfile, buf)
-    # TODO(johnnychen94): load back with `JpegTurbo.decode`
-    return ImageMagick.load(tmpfile)
+    return jpeg_decode(tmpfile)
 end
 
 @testset "JpegTurbo.jl" begin
@@ -47,6 +46,7 @@ end
     end
 
     include("tst_encode.jl")
+    include("tst_decode.jl")
     if Threads.nthreads() > 1
         @info "Multi-threads test: enabled"
         include("tst_multithreads.jl")

test/tst_decode.jl

@@ -0,0 +1,49 @@
+@testset "jpeg_decode" begin
+    img_rgb = testimage("lighthouse")
+
+    tmpfile = joinpath(tmpdir, "tmp.jpg")
+    jpeg_encode(tmpfile, img_rgb)
+
+    data = jpeg_decode(tmpfile)
+
+    @test jpeg_decode(tmpfile; transpose=true) == data'
+
+    # ensure default keyword values are not changed by accident
+    @test jpeg_decode(tmpfile) ≈
+        jpeg_decode(RGB, tmpfile; transpose=false, scale_ratio=1) ≈
+        jpeg_decode(tmpfile; transpose=false, scale_ratio=1)
+
+
+    # TODO(johnnychen94): support IO and in-memory buffer
+    @test_broken jpeg_decode(jpeg_encode(img_rgb))
+    @test_broken open(jpeg_decode, tmpfile, "r")
+
+    @testset "colorspace" begin
+        native_color_spaces = [Gray, RGB, BGR, RGBA, BGRA, ABGR, ARGB]
+        ext_color_spaces = [YCbCr, RGBX, XRGB, Lab, YIQ] # supported by Colors.jl
+        for CT in [native_color_spaces..., ext_color_spaces...]
+            data = jpeg_decode(CT, tmpfile)
+            @test eltype(data) <: CT
+            if CT == Gray
+                @test assess_psnr(data, Gray.(img_rgb)) > 34.92
+            else
+                @test assess_psnr(RGB.(data), img_rgb) > 33.87
+            end
+        end
+    end
+
+    @testset "scale_ratio" begin
+        data = jpeg_decode(tmpfile; scale_ratio=0.25)
+        @test size(data) == (128, 192) == 0.25 .* size(img_rgb)
+
+        # `jpeg_decode` will map input `scale_ratio` to allowed values.
+        data = jpeg_decode(tmpfile; scale_ratio=0.3)
+        @test size(data) == (128, 192) != 0.3 .* size(img_rgb)
+    end
+
+    @testset "transpose" begin
+        jpeg_encode(tmpfile, img_rgb; transpose=true)
+        data = jpeg_decode(tmpfile; transpose=true)
+        @test assess_psnr(data, img_rgb) > 33.95
+    end
+end

test/tst_encode.jl

@@ -6,6 +6,8 @@ img_rgb = testimage("lighthouse")
     for CT in [Gray, RGB, #=YCbCr,=# #=RGBX,=# BGR, #=XRGB,=# RGBA, BGRA, ABGR, ARGB]
         img = CT.(img_rgb)
         data = decode_encode(img)
+        @test eltype(data) <: Union{Gray, RGB}
+        @test size(data) == size(img)
         @test data ≈ decode_encode(float32.(img))
 
         # ensure default keyword values are not changed by accident
@@ -17,12 +19,12 @@ end
 # keyword checks
 @testset "quality" begin
     img = testimage("cameraman")
-    psnr_refs = Dict(
-        1   => 24.5647,
-        10  => 31.3434,
-        50  => 38.8765,
-        100 => 59.3770,
-    )
+    psnr_refs = [
+        1   => 24.63,
+        10  => 31.34,
+        50  => 38.87,
+        100 => 59.31,
+    ]
     for (q, r) in psnr_refs
         v = assess_psnr(img, decode_encode(img, quality=q))
         @test v >= r

test/tst_multithreads.jl

@@ -11,4 +11,18 @@
         end
         @test all(out .== Ref(ref))
     end
+
+    @testset "jpeg_decode" begin
+        out = [similar(img) for _ in 1:Threads.nthreads()]
+        tmpdir = tempdir()
+
+        for i in 1:Threads.nthreads()
+            jpeg_encode(joinpath(tmpdir, "$i.jpg"), img)
+        end
+        ref = jpeg_decode(joinpath(tmpdir, "1.jpg"))
+        Threads.@threads for i in 1:Threads.nthreads()
+            out[i] = jpeg_decode(joinpath(tmpdir, "$i.jpg"))
+        end
+        @test all(out .== Ref(ref))
+    end
 end