using Random, Images function get_images_and_labels(root_dir::String, image_size::Tuple{Int,Int}, shuffle_data::Bool, val_split::Float64, seed::Int=42) # initialise storage arrays label_categ = String[] image_paths = String[] labels = String[] # get image paths and labels for (index, (root, dirs, files)) in enumerate(walkdir(root_dir)) (root == root_dir) && (label_categ = dirs) if index > 1 new_images = map(x -> (root_dir * label_categ[index-1] * '/' * x), files) new_labels = repeat([label_categ[index-1]], length(files)) image_paths = vcat(image_paths, new_images) labels = vcat(labels, new_labels) end end # load and scale images imgs = Images.load.(image_paths) # note: image_size is expected in the format (height, width) imgs = map(image -> Images.imresize(image, image_size), imgs) all_data = hcat(imgs, labels) # shuffle data Random.Random.seed!(seed) shuffle_data && (all_data = all_data[Random.shuffle(begin:end),:]) # train / val split split_index = trunc(Int, (size(all_data)[1] * val_split)) train = all_data[split_index+1:end,:] val = all_data[begin:split_index,:] return train[:,begin], train[:,end], val[:,begin], val[:,end], label_categ end
Hosted onDeepnote