Add img_size argument

JuliaML · Feb 8, 2023 · 09d5be4 · 09d5be4
1 parent 95b13d9
commit 09d5be4
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/src/datasets/vision/imagenet.jl b/src/datasets/vision/imagenet.jl
@@ -92,6 +92,7 @@ Dict{String, Any} with 8 entries:
   "class_names"       => Vector{SubString{String}}[["tench", "Tinca tinca"], ["goldfish", "C
   "metadata_path"     => "/Users/funks/.julia/datadeps/ImageNet/devkit/data/meta.mat"
   "n_classes"         => 1000
+  "img_size"          => (224, 224)
   "wnid_to_label"     => Dict("n07693725"=>932, "n03775546"=>660, "n01689811"=>45, "n0210087
 
 julia> dataset.metadata["class_names"][y]
@@ -120,6 +121,7 @@ ImageNet(Tx::Type; kws...) = ImageNet(; Tx, kws...)
 function ImageNet(
     Tx::Type,
     split::Symbol;
+    img_size::Tuple{Int,Int}=(224, 224),
     preprocess=ImageNetReader.default_preprocess,
     inverse_preprocess=ImageNetReader.default_inverse_preprocess,
     dir=nothing,
@@ -157,9 +159,10 @@ function ImageNet(
     metadata["features_dir"] = features_dir
     metadata["n_observations"] = n_observations
     metadata["n_classes"] = ImageNetReader.NCLASSES
+    metadata["img_size"] = img_size
 
     # Create FileDataset
-    dataset = ImageNetReader.get_file_dataset(Tx, preprocess, features_dir)
+    dataset = ImageNetReader.get_file_dataset(Tx, img_size, preprocess, features_dir)
     @assert length(dataset) == n_observations
 
     targets = [

diff --git a/src/datasets/vision/imagenet_reader/ImageNetReader.jl b/src/datasets/vision/imagenet_reader/ImageNetReader.jl
@@ -8,16 +8,17 @@ import ..@lazy
 @lazy import JpegTurbo = "b835a17e-a41a-41e7-81f0-2f016b05efe0"
 
 const NCLASSES = 1000
-const IMGSIZE = (224, 224)
 
 include("preprocess.jl")
 
-function get_file_dataset(Tx::Type{<:Real}, preprocess::Function, dir::AbstractString)
+function get_file_dataset(
+    Tx::Type{<:Real}, img_size::Tuple{Int,Int}, preprocess::Function, dir::AbstractString
+)
     # Construct a function that loads images from FileDataset path,
     # applies preprocessing and converts to type Tx.
     function load_image(file::AbstractString)
-        im = JpegTurbo.jpeg_decode(RGB{Tx}, file; preferred_size=IMGSIZE)
-        return Tx.(preprocess(im))
+        im = JpegTurbo.jpeg_decode(RGB{Tx}, file; preferred_size=img_size)
+        return preprocess(im, img_size)
     end
     return FileDataset(load_image, dir, "*.JPEG")
 end

diff --git a/src/datasets/vision/imagenet_reader/preprocess.jl b/src/datasets/vision/imagenet_reader/preprocess.jl
@@ -1,7 +1,7 @@
 # Image preprocessing defaults for ImageNet models.
 
-function default_preprocess(im::AbstractMatrix{<:AbstractRGB})
-    im = channelview(center_crop(im))
+function default_preprocess(im::AbstractMatrix{<:AbstractRGB}, outsize)
+    im = channelview(center_crop(im, outsize))
     return PermutedDimsArray(im, (3, 2, 1)) # Convert from Image.jl's CHW to Flux's WHC
 end
 
@@ -12,7 +12,7 @@ function default_inverse_preprocess(x::AbstractArray{T,N}) where {T,N}
 end
 
 # Take rectangle of pixels of shape `outsize` at the center of image `im`
-function center_crop(im::AbstractMatrix, outsize=IMGSIZE)
+function center_crop(im::AbstractMatrix, outsize)
     h2, w2 = div.(outsize, 2) # half height, half width of view
     h_adjust, w_adjust = _adjust.(outsize)
     return @view im[