parallelize decoding (#127)

* parallelize decoding * use _read(::IO, ::AbstractArray) for backwards compatibility --------- Co-authored-by: Tamas Nagy <[email protected]>
tlnagy · Nov 28, 2023 · 29c5038 · 29c5038 · tlnagy · Dec 6, 2023
1 parent 4613107
commit 29c5038
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 18 deletions.
diff --git a/src/compression.jl b/src/compression.jl
@@ -20,13 +20,13 @@ function Base.read!(tfs::TiffFileStrip, arr::AbstractArray{T, N}, ::Val{COMPRESS
     nxt = Array{UInt8}(undef, 1)
     arr = reinterpret(UInt8, arr)
     while pos < length(arr)
-        read!(tfs.tf, nbit)
+        read!(tfs.io, nbit)
         n = nbit[1]
         if 0 <= n <= 127
-            read!(tfs.tf, view(arr, pos:pos+n))
+            _read!(tfs.io, view(arr, pos:pos+n))
             pos += n
         elseif -127 <= n <= -1
-            read!(tfs.tf, nxt)
+            read!(tfs.io, nxt)
             arr[pos:(pos-n)] .= nxt[1]
             pos += -n
         end
@@ -35,11 +35,11 @@ function Base.read!(tfs::TiffFileStrip, arr::AbstractArray{T, N}, ::Val{COMPRESS
 end
 
 function Base.read!(tfs::TiffFileStrip, arr::AbstractArray, ::Val{COMPRESSION_DEFLATE})
-    readbytes!(InflateZlibStream(tfs.tf.io.io), reinterpret(UInt8, vec(arr)))
+    readbytes!(InflateZlibStream(tfs.io), reinterpret(UInt8, vec(arr)))
 end
 
 function Base.read!(tfs::TiffFileStrip, arr::AbstractArray, ::Val{COMPRESSION_ADOBE_DEFLATE})
-    readbytes!(InflateZlibStream(tfs.tf.io.io), reinterpret(UInt8, vec(arr)))
+    readbytes!(InflateZlibStream(tfs.io), reinterpret(UInt8, vec(arr)))
 end
 
 function lzw_decode!(io, arr::AbstractArray)

diff --git a/src/ifds.jl b/src/ifds.jl
@@ -212,17 +212,16 @@ pixels in the image
 
 $(FIELDS)
 """
-struct TiffFileStrip{O, S, P}
-    """The file stream"""
-    tf::TiffFile{O, S}
+struct TiffFileStrip{O, T}
+    """Strip data"""
+    io::IOBuffer
+
     """The IFD corresponding to this strip"""
     ifd::IFD{O}
-    """The number of bytes in this strip"""
-    bytes::Int
 end
 
-Base.read!(tfs::TiffFileStrip, arr::AbstractArray) = read!(tfs.tf, arr)
-Base.bytesavailable(tfs::TiffFileStrip) = tfs.bytes
+Base.read!(tfs::TiffFileStrip, arr::AbstractArray) = _read!(tfs.io, arr)
+Base.bytesavailable(tfs::TiffFileStrip) = bytesavailable(tfs.io)
 
 function Base.read!(target::AbstractArray{T, N}, tf::TiffFile{O, S}, ifd::IFD{O}) where {T, N, O, S}
     if PLANARCONFIG in ifd
@@ -237,11 +236,12 @@ function Base.read!(target::AbstractArray{T, N}, tf::TiffFile{O, S}, ifd::IFD{O}
         # number of input bytes in each strip or tile
         encoded_bytes = istiled(ifd) ? ifd[TILEBYTECOUNTS].data : ifd[STRIPBYTECOUNTS].data
 
+        rows = nrows(ifd)
+        cols = ncols(ifd)
+
         if istiled(ifd)
             strip_pixels = map(_ -> tilecols(ifd) * tilerows(ifd), encoded_bytes)
         else
-            rows = nrows(ifd)
-            cols = ncols(ifd)
             nstrips = length(encoded_bytes)
             rowsperstrip = getdata(Int, ifd, ROWSPERSTRIP, rows)
 
@@ -251,17 +251,37 @@ function Base.read!(target::AbstractArray{T, N}, tf::TiffFile{O, S}, ifd::IFD{O}
             @assert sum(strip_pixels) == rows * cols
         end
 
+        parallel_enabled = something(tryparse(Bool, get(ENV, "JULIA_IMAGES_PARALLEL", "1")), false)
+        do_parallel = parallel_enabled && rows * cols > 250_000 # pixels
+
         start = 1
         comp = Val(compression)
         rtype = rawtype(ifd)
+        tasks::Vector{Task} = []
         for (offset, len, bytes) in zip(offsets, strip_pixels, encoded_bytes)
             seek(tf, offset)
-            tfs = TiffFileStrip{O, S, rtype}(tf, ifd, bytes)
             arr = view(target, start:(start+len-1))
-            read!(tfs, arr, comp)
-            reverse_prediction!(tfs, arr)
+            data = Vector{UInt8}(undef, bytes)
+            read!(tf, data)
+            tfs = TiffFileStrip{O, rtype}(IOBuffer(data), ifd)
+
+            function go(tfs, arr, comp)
+                read!(tfs, arr, comp)
+                reverse_prediction!(tfs, arr)
+            end
+
+            if do_parallel
+                push!(tasks, Threads.@spawn go(tfs, arr, comp))
+            else
+                go(tfs, arr, comp)
+            end
+
             start += len
         end
+
+        for task in tasks
+            wait(task)
+        end
     else
         seek(tf, first(offsets))
         read!(tf, target, compression)
@@ -324,7 +344,7 @@ function Base.write(tf::TiffFile{O}, ifd::IFD{O}) where {O <: Unsigned}
     return ifd_end_pos
 end
 
-function reverse_prediction!(tfs::TiffFileStrip{O, S, P}, arr::AbstractArray{T, N}) where {O, S, P, T, N}
+function reverse_prediction!(tfs::TiffFileStrip{O, P}, arr::AbstractArray{T,N}) where {O, P, T, N}
     pred::Int = predictor(tfs.ifd)
     spp::Int = nsamples(tfs.ifd)
     if pred == 2

diff --git a/src/load.jl b/src/load.jl
@@ -6,6 +6,9 @@ true, false, and false by default, respectively. Setting `verbose` to false
 will hide the loading bar, while setting either `lazyio` or `mmap` to true
 defer loading until the data are needed (by either of two mechanisms).
 
+Parallelism is enabled by default, but can be disabled by setting
+`JULIA_IMAGES_PARALLEL`=false in your environment
+
 See [Lazy TIFFs](@ref) for more details about memory-mapping and lazy I/O.
 """
 function load(filepath::String; mode = "r", kwargs...)