feat(decode): support IO and in-memory data (#13)

johnnychen94 · web-flow · commit abd540e04aeb · 2022-02-04T04:02:58.000+08:00
diff --git a/README.md b/README.md
@@ -26,14 +26,16 @@ jpeg_encode(img; kwargs...) -> Vector{UInt8}
 
 ```julia
 jpeg_decode([T,] filename::AbstractString; kwargs...) -> Matrix{T}
+jpeg_decode([T,] io::IO; kwargs...) -> Matrix{T}
+jpeg_decode([T,] data::Vector{UInt8}; kwargs...) -> Matrix{T}
 ```
 
 ## Feature set
 
 | function             | filename | IOStream | in-memory buffer     | pre-allocated output | multi-threads |
 | -------------------- | -------- | -------- | -------------------- | -------------------  | ------------- |
 | `jpeg_encode`        | x        | x        | x                    |                      | x             |
-| `jpeg_decode`        | x        |          |                      |                      | x             |
+| `jpeg_decode`        | x        | x        | x                    |                      | x             |
 | `ImageMagick.save`   | x        | x        | x                    |                      | x             |
 | `ImageMagick.load`   | x        | x        | x                    |                      | x             |
 | `QuartzImageIO.save` | x        | x        | x (`FileIO.Stream`)  |                      | x             |
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -21,14 +21,16 @@ jpeg_encode(img; kwargs...) -> Vector{UInt8}
 
 ```julia
 jpeg_decode([T,] filename::AbstractString; kwargs...) -> Matrix{T}
+jpeg_decode([T,] io::IO; kwargs...) -> Matrix{T}
+jpeg_decode([T,] data::Vector{UInt8}; kwargs...) -> Matrix{T}
 ```
 
 ## Feature set
 
 | function             | filename | IOStream | in-memory buffer     | pre-allocated output | multi-threads |
 | -------------------- | -------- | -------- | -------------------- | -------------------  | ------------- |
 | `jpeg_encode`        | x        | x        | x                    |                      | x             |
-| `jpeg_decode`        | x        |          |                      |                      | x             |
+| `jpeg_decode`        | x        | x        | x                    |                      | x             |
 | `ImageMagick.save`   | x        | x        | x                    |                      | x             |
 | `ImageMagick.load`   | x        | x        | x                    |                      | x             |
 | `QuartzImageIO.save` | x        | x        | x (`FileIO.Stream`)  |                      | x             |
diff --git a/src/decode.jl b/src/decode.jl
@@ -1,7 +1,10 @@
 """
     jpeg_decode([T,] filename::AbstractString; kwargs...) -> Matrix{T}
+    jpeg_decode([T,] io::IO; kwargs...) -> Matrix{T}
+    jpeg_decode([T,] data::Vector{UInt8}; kwargs...) -> Matrix{T}
 
-Decode the JPEG image from given I/O stream as colorant matrix.
+Decode the JPEG image as colorant matrix. The source data can be either a filename, an IO
+, or an in-memory byte sequence.
 
 # parameters
 
@@ -47,11 +50,10 @@ filename = testimage("earth", download_only=true)
 """
 function jpeg_decode(
         ::Type{CT},
-        filename::AbstractString;
+        data::Vector{UInt8};
         transpose=false,
         scale_ratio=1) where CT<:Colorant
-    infile = ccall(:fopen, Libc.FILE, (Cstring, Cstring), filename, "rb")
-    @assert infile.ptr != C_NULL
+    _jpeg_check_bytes(data)
     out_CT, jpeg_cls = _jpeg_out_color_space(CT)
 
     cinfo_ref = Ref(LibJpeg.jpeg_decompress_struct())
@@ -60,7 +62,7 @@ function jpeg_decode(
         cinfo = cinfo_ref[]
         cinfo.err = LibJpeg.jpeg_std_error(jerr)
         LibJpeg.jpeg_create_decompress(cinfo_ref)
-        LibJpeg.jpeg_stdio_src(cinfo_ref, infile)
+        LibJpeg.jpeg_mem_src(cinfo_ref, data, length(data))
         LibJpeg.jpeg_read_header(cinfo_ref, true)
 
         # set decompression parameters, if given
@@ -87,13 +89,21 @@ function jpeg_decode(
         end
     finally
         LibJpeg.jpeg_destroy_decompress(cinfo_ref)
-        ccall(:fclose, Cint, (Ptr{Libc.FILE},), infile)
     end
 end
-function jpeg_decode(filename::AbstractString; kwargs...)
-    return jpeg_decode(_default_out_color_space(filename), filename; kwargs...)
+jpeg_decode(data; kwargs...) = jpeg_decode(_default_out_color_space(data), data; kwargs...)
+
+# TODO(johnnychen94): support Progressive JPEG
+# TODO(johnnychen94): support partial decoding
+function jpeg_decode(::Type{CT}, filename::AbstractString; kwargs...) where CT<:Colorant
+    open(filename, "r") do io
+        jpeg_decode(CT, io; kwargs...)
+    end
 end
 
+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!(out::Matrix{<:Colorant}, cinfo_ref::Ref{LibJpeg.jpeg_decompress_struct})
     row_stride = size(out, 1) * length(eltype(out))
     buf = Vector{UInt8}(undef, row_stride)
@@ -120,6 +130,7 @@ 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)
+    _jpeg_check_bytes(filename)
     infile = ccall(:fopen, Libc.FILE, (Cstring, Cstring), filename, "rb")
     @assert infile.ptr != C_NULL
     cinfo_ref = Ref(LibJpeg.jpeg_decompress_struct())
@@ -137,6 +148,22 @@ function _default_out_color_space(filename::AbstractString)
     end
 end
 
+function _default_out_color_space(data::Vector{UInt8})
+    _jpeg_check_bytes(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)
+        LibJpeg.jpeg_mem_src(cinfo_ref, data, length(data))
+        LibJpeg.jpeg_read_header(cinfo_ref, true)
+        LibJpeg.jpeg_calc_output_dimensions(cinfo_ref)
+        return jpeg_color_space(cinfo_ref[].out_color_space)
+    finally
+        LibJpeg.jpeg_destroy_decompress(cinfo_ref)
+    end
+end
+
 function _jpeg_out_color_space(::Type{CT}) where CT
     try
         n0f8(CT), jpeg_color_space(n0f8(CT))
@@ -145,3 +172,25 @@ function _jpeg_out_color_space(::Type{CT}) where CT
         RGB{N0f8}, jpeg_color_space(RGB{N0f8})
     end
 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)
+    seekend(io)
+    nbytes = position(io)
+    nbytes > 623 || throw(ArgumentError("Invalid number of bytes."))
+
+    buf = UInt8[]
+    seekstart(io)
+    readbytes!(io, buf, 623)
+    seek(io, nbytes-2)
+    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."))
+    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
diff --git a/test/Project.toml b/test/Project.toml
@@ -3,7 +3,6 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
-ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
 ImageQualityIndexes = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990"
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -5,19 +5,12 @@ using Aqua
 using Documenter
 using TestImages
 using ImageQualityIndexes
-using ImageMagick
 using ImageCore
 
 # ensure TestImages artifacts are downloaded before running documenter test
 testimage("cameraman")
 
-tmpdir = tempdir()
-function decode_encode(img; kwargs...)
-    tmpfile = joinpath(tmpdir, "tmp.jpg")
-    buf = @inferred jpeg_encode(img; kwargs...)
-    write(tmpfile, buf)
-    return jpeg_decode(tmpfile)
-end
+const tmpdir = tempdir()
 
 @testset "JpegTurbo.jl" begin
     if !Sys.iswindows() # DEBUG
diff --git a/test/tst_decode.jl b/test/tst_decode.jl
@@ -1,28 +1,38 @@
 @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'
+    img_rgb_bytes = jpeg_encode(img_rgb)
 
     # 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)
+    @test jpeg_decode(img_rgb_bytes) ≈
+        jpeg_decode(RGB, img_rgb_bytes; transpose=false, scale_ratio=1) ≈
+        jpeg_decode(img_rgb_bytes; transpose=false, scale_ratio=1)
+
+    @testset "filename and IOStream" begin
+        tmpfile = joinpath(tmpdir, "tmp.jpg")
+        jpeg_encode(tmpfile, img_rgb)
+        @test read(tmpfile) == img_rgb_bytes
+
+        # IOStream
+        img = open(tmpfile, "r") do io
+            jpeg_decode(io)
+        end
+        @test img == jpeg_decode(img_rgb_bytes)
 
+        img = open(tmpfile, "r") do io
+            jpeg_decode(Gray, io; scale_ratio=0.5)
+        end
+        @test img == jpeg_decode(Gray, img_rgb_bytes; scale_ratio=0.5)
 
-    # TODO(johnnychen94): support IO and in-memory buffer
-    @test_broken jpeg_decode(jpeg_encode(img_rgb))
-    @test_broken open(jpeg_decode, tmpfile, "r")
+        # filename
+        @test jpeg_decode(tmpfile) == jpeg_decode(img_rgb_bytes)
+        @test jpeg_decode(Gray, tmpfile; scale_ratio=0.5) == jpeg_decode(Gray, img_rgb_bytes; scale_ratio=0.5)
+    end
 
     @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)
+            data = jpeg_decode(CT, img_rgb_bytes)
             @test eltype(data) <: CT
             if CT == Gray
                 @test assess_psnr(data, Gray.(img_rgb)) > 34.92
@@ -33,17 +43,21 @@
     end
 
     @testset "scale_ratio" begin
-        data = jpeg_decode(tmpfile; scale_ratio=0.25)
+        data = jpeg_decode(img_rgb_bytes; 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)
+        data = jpeg_decode(img_rgb_bytes; 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)
+        data = jpeg_decode(jpeg_encode(img_rgb; transpose=true); transpose=true)
         @test assess_psnr(data, img_rgb) > 33.95
     end
+
+    @testset "integrity check" begin
+        @test_throws ArgumentError jpeg_decode(UInt8[])
+        @test_throws ArgumentError jpeg_decode(img_rgb_bytes[1:600])
+    end
 end
diff --git a/test/tst_encode.jl b/test/tst_encode.jl
@@ -5,14 +5,15 @@ img_rgb = testimage("lighthouse")
 @testset "basic" begin
     for CT in [Gray, RGB, #=YCbCr,=# #=RGBX,=# BGR, #=XRGB,=# RGBA, BGRA, ABGR, ARGB]
         img = CT.(img_rgb)
-        data = decode_encode(img)
+        data = jpeg_decode(jpeg_encode(img))
         @test eltype(data) <: Union{Gray, RGB}
         @test size(data) == size(img)
-        @test data ≈ decode_encode(float32.(img))
+        @test data ≈ jpeg_decode(jpeg_encode(float32.(img)))
 
         # ensure default keyword values are not changed by accident
-        @test data == decode_encode(img, transpose=false)
-        @test decode_encode(img, transpose=true) == decode_encode(img', transpose=false)
+        @test data == jpeg_decode(jpeg_encode(img, transpose=false))
+        @test jpeg_decode(jpeg_encode(img, transpose=true)) ==
+            jpeg_decode(jpeg_encode(img', transpose=false))
     end
 
     # numerical array is treated as Gray image
@@ -35,7 +36,7 @@ end
         100 => 59.31,
     ]
     for (q, r) in psnr_refs
-        v = assess_psnr(img, decode_encode(img, quality=q))
+        v = assess_psnr(img, jpeg_decode(jpeg_encode(img, quality=q)))
         @test v >= r
     end
 end
