Skip to content

Instantly share code, notes, and snippets.

@TugdualKerjan

TugdualKerjan/Cargo.toml

Last active July 1, 2025 10:19
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save TugdualKerjan/760ba29edd6a664b810307227f3b5424 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save TugdualKerjan/760ba29edd6a664b810307227f3b5424 to your computer and use it in GitHub Desktop.
Download ZIP
ImageNetResize with rust
Raw
Cargo.toml
[package]
name = "image_resizer"
version = "0.1.0"
edition = "2021"
[dependencies]
image = "0.24"
rayon = "1.7"
clap = "3.2"
Raw
Explanation.md

Compile and run the thing like this

cargo build --release
./target/release/image_resizer -r /mnt/data/ILSVRC2012_raw/val -o /mnt/data/ILSVRC2012_32/val/ -r -s 32 -a bicubic
Raw
main.rs
use std::path::{Path, PathBuf};
use std::fs;
use std::sync::{Arc, Mutex};
use std::io::Write;
use image::{ImageFormat, GenericImageView, imageops::FilterType};
use rayon::prelude::*;
use clap::{Arg, App};
use std::collections::HashMap;
// Define a mapping for resampling algorithms
fn get_filter_type_map() -> HashMap<String, FilterType> {
let mut map = HashMap::new();
map.insert("lanczos".to_string(), FilterType::Lanczos3);
map.insert("nearest".to_string(), FilterType::Nearest);
map.insert("bilinear".to_string(), FilterType::Triangle);
map.insert("bicubic".to_string(), FilterType::CatmullRom);
// Box filter is equivalent to FilterType::Nearest in the current image version
map.insert("box".to_string(), FilterType::Nearest);
// Hamming is not available in current image version, use Triangle instead
map.insert("hamming".to_string(), FilterType::Triangle);
map
}
// Resize a single image
fn resize_image(
input_path: &Path,
output_path: &Path,
size: u32,
filter: FilterType,
error_log: Arc<Mutex<fs::File>>,
) -> Result<(), String> {
// Create parent directories if they don't exist
if let Some(parent) = output_path.parent() {
fs::create_dir_all(parent).map_err(|e| format!("Failed to create directory: {}", e))?;
}
// Attempt to open and process the image
match image::open(input_path) {
Ok(img) => {
// Convert to RGB if needed
let img_rgb = if img.color().has_alpha() {
image::DynamicImage::ImageRgba8(img.to_rgba8())
} else if img.color().channel_count() == 1 {
image::DynamicImage::ImageRgb8(img.to_rgb8())
} else {
img
};
// Get original dimensions
let (width, height) = img_rgb.dimensions();
// Calculate dimensions for resizing the largest side to target size
let (new_width, new_height) = if width < height {
let ratio = size as f32 / width as f32;
(size, (height as f32 * ratio).round() as u32)
} else {
let ratio = size as f32 / height as f32;
((width as f32 * ratio).round() as u32, size)
};
// Resize the image with the largest dimension at target size
let mut resized = img_rgb.resize(new_width, new_height, filter);
// Now crop to square from the center
let final_image = if new_width != new_height {
// Calculate crop coordinates (center crop)
let x = if new_width > size { (new_width - size) / 2 } else { 0 };
let y = if new_height > size { (new_height - size) / 2 } else { 0 };
let crop_width = std::cmp::min(size, new_width);
let crop_height = std::cmp::min(size, new_height);
// Crop to square
resized.crop(x, y, crop_width, crop_height)
} else {
// Already square at target size
resized
};
// Save the final image
final_image.save_with_format(output_path, ImageFormat::Png)
.map_err(|e| format!("Failed to save image: {}", e))?;
Ok(())
},
Err(e) => {
// Log the error
let error_msg = format!("Couldn't resize: {} - Error: {}\n", input_path.display(), e);
if let Ok(mut log) = error_log.lock() {
let _ = log.write_all(error_msg.as_bytes());
}
Err(error_msg)
}
}
}
// Process a directory of images
fn process_directory(
input_dir: &Path,
output_dir: &Path,
size: u32,
filter: FilterType,
error_log: Arc<Mutex<fs::File>>,
) -> Result<(), String> {
// Ensure the output directory exists
fs::create_dir_all(output_dir).map_err(|e| format!("Failed to create output directory: {}", e))?;
// Get all files in the directory
let entries = match fs::read_dir(input_dir) {
Ok(entries) => entries,
Err(e) => return Err(format!("Failed to read directory {}: {}", input_dir.display(), e)),
};
// Process each file in parallel
let results: Vec<Result<(), String>> = entries
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().is_file())
.collect::<Vec<_>>()
.par_iter()
.map(|entry| {
let input_path = entry.path();
let file_stem = input_path.file_stem().unwrap_or_default();
let output_path = output_dir.join(file_stem).with_extension("png");
resize_image(&input_path, &output_path, size, filter, Arc::clone(&error_log))
})
.collect();
// Check for errors
let errors: Vec<String> = results.into_iter()
.filter_map(|result| result.err())
.collect();
if errors.is_empty() {
Ok(())
} else {
Err(format!("Errors occurred while processing directory: {:?}", errors))
}
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Create the CLI using clap
let matches = App::new("Image Resizer")
.version("1.0")
.about("Resizes images in a directory")
.arg(Arg::with_name("in_dir")
.short('i')
.long("in_dir")
.value_name("INPUT_DIR")
.help("Input directory with source images")
.required(true)
.takes_value(true))
.arg(Arg::with_name("out_dir")
.short('o')
.long("out_dir")
.value_name("OUTPUT_DIR")
.help("Output directory for resized images")
.required(true)
.takes_value(true))
.arg(Arg::with_name("size")
.short('s')
.long("size")
.value_name("SIZE")
.help("Size of an output image (e.g. 32 results in (32x32) image)")
.default_value("128")
.takes_value(true))
.arg(Arg::with_name("algorithm")
.short('a')
.long("algorithm")
.value_name("ALGORITHM")
.help("Algorithm used for resampling: lanczos, nearest, bilinear, bicubic, box, hamming")
.default_value("box")
.takes_value(true))
.arg(Arg::with_name("recurrent")
.short('r')
.long("recurrent")
.help("Process all subfolders in this folder (1 lvl deep)")
.takes_value(false))
.arg(Arg::with_name("full")
.short('f')
.long("full")
.help("Use all algorithms, create subdirectory for each algorithm output")
.takes_value(false))
.arg(Arg::with_name("every_nth")
.short('e')
.long("every_nth")
.value_name("N")
.help("Use if you don't want to take all classes, if -e 10 then takes every 10th class")
.default_value("1")
.takes_value(true))
.get_matches();
// Parse arguments
let in_dir = matches.value_of("in_dir").unwrap();
let out_dir = matches.value_of("out_dir").unwrap();
let size: u32 = matches.value_of("size").unwrap().parse()?;
let alg_str = matches.value_of("algorithm").unwrap().to_lowercase();
let recurrent = matches.is_present("recurrent");
let full = matches.is_present("full");
let every_nth: usize = matches.value_of("every_nth").unwrap().parse()?;
println!("Starting...");
// Get filter types map
let filter_map = get_filter_type_map();
// Determine which algorithms to use
let algs = if full {
filter_map.keys().cloned().collect::<Vec<String>>()
} else {
vec![alg_str.clone()]
};
// Create error log file
let error_log = Arc::new(Mutex::new(
fs::OpenOptions::new()
.create(true)
.append(true)
.open("log.txt")?
));
// Process all algorithms
for alg in algs {
println!("Using algorithm {}...", alg);
let filter = match filter_map.get(&alg) {
Some(filter) => *filter,
None => {
println!("Sorry but this algorithm ({}) is not available, use help for more info.", alg);
continue;
}
};
// let current_out_dir = PathBuf::from(out_dir).join(&alg);
let current_out_dir = PathBuf::from(out_dir);
if recurrent {
println!("Recurrent for all folders in in_dir: {}", in_dir);
// Get all subdirectories
let folders: Vec<PathBuf> = fs::read_dir(in_dir)?
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().is_dir())
.map(|entry| entry.path())
.collect();
// Process subdirectories in parallel with every_nth filter
folders.into_par_iter()
.enumerate()
.filter(|(i, _)| i % every_nth == 0)
.for_each(|(_, folder)| {
let folder_name = folder.file_name().unwrap_or_default();
let input_path = PathBuf::from(in_dir).join(folder_name);
let output_path = current_out_dir.join(folder_name);
println!("Processing folder: {}", input_path.display());
if let Err(e) = process_directory(
&input_path,
&output_path,
size,
filter,
Arc::clone(&error_log)
) {
println!("Error processing folder {}: {}", input_path.display(), e);
}
});
} else {
println!("For folder {}", in_dir);
if let Err(e) = process_directory(
&PathBuf::from(in_dir),
&current_out_dir,
size,
filter,
Arc::clone(&error_log)
) {
println!("Error: {}", e);
}
}
}
println!("Finished.");
Ok(())
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment