From a5928d3b0e29b5a55e1296dc4e19ec65589660d5 Mon Sep 17 00:00:00 2001
From: Robert Knight <robertknight@gmail.com>
Date: Wed, 18 Dec 2024 18:07:43 +0000
Subject: [PATCH] Add EDSR image super resolution example

---
 rten-examples/Cargo.toml         |  5 ++
 rten-examples/README.md          |  1 +
 rten-examples/src/edsr.rs        | 85 ++++++++++++++++++++++++++++++++
 rten-examples/src/export-edsr.py | 28 +++++++++++
 4 files changed, 119 insertions(+)
 create mode 100644 rten-examples/src/edsr.rs
 create mode 100644 rten-examples/src/export-edsr.py

diff --git a/rten-examples/Cargo.toml b/rten-examples/Cargo.toml
index 8fbfec90..4fccd467 100644
--- a/rten-examples/Cargo.toml
+++ b/rten-examples/Cargo.toml
@@ -57,6 +57,11 @@ name = "distilvit"
 path = "src/distilvit.rs"
 test = false
 
+[[bin]]
+name = "edsr"
+path = "src/edsr.rs"
+test = false
+
 [[bin]]
 name = "imagenet"
 path = "src/imagenet.rs"
diff --git a/rten-examples/README.md b/rten-examples/README.md
index 3b40bfeb..d8bd5791 100644
--- a/rten-examples/README.md
+++ b/rten-examples/README.md
@@ -64,6 +64,7 @@ The examples have been chosen to cover common tasks and popular models.
 - **depth_anything** - Monocular depth estimation using [Depth Anything](https://github.com/LiheYoung/Depth-Anything)
 - **detr** - Object detection using [DETR](https://research.facebook.com/publications/end-to-end-object-detection-with-transformers/)
 - **distilvit** - Image captioning using [Mozilla's DistilViT](https://hacks.mozilla.org/2024/05/experimenting-with-local-alt-text-generation-in-firefox-nightly/)
+- **edsr** - Image super-resolution using [EDSR](https://github.com/eugenesiow/super-image)
 - **nougat** - Extract text from academic PDFs as Markdown using [Nougat](https://github.com/facebookresearch/nougat/)
 - **rmbg** - Background removal using [BRIA Background Removal](https://huggingface.co/briaai/RMBG-1.4)
 - **segment_anything** - Image segmentation using [Segment Anything](https://segment-anything.com)
diff --git a/rten-examples/src/edsr.rs b/rten-examples/src/edsr.rs
new file mode 100644
index 00000000..84083696
--- /dev/null
+++ b/rten-examples/src/edsr.rs
@@ -0,0 +1,85 @@
+use std::collections::VecDeque;
+use std::error::Error;
+
+use rten::Model;
+use rten_imageio::{read_image, write_image};
+use rten_tensor::prelude::*;
+use rten_tensor::{NdTensor, Tensor};
+
+struct Args {
+    model: String,
+    image: String,
+    output: String,
+}
+
+fn parse_args() -> Result<Args, lexopt::Error> {
+    use lexopt::prelude::*;
+
+    let mut values = VecDeque::new();
+    let mut parser = lexopt::Parser::from_env();
+
+    while let Some(arg) = parser.next()? {
+        match arg {
+            Value(val) => values.push_back(val.string()?),
+            Long("help") => {
+                println!(
+                    "Perform super-resolution (upscaling) on an image
+
+Usage: {bin_name} <model> <image> [<output>]
+
+Args:
+
+  <model> - Super-resolution model
+  <image> - Image to process
+  <output> - Path to save image to. Defaults to \"output.png\".
+",
+                    bin_name = parser.bin_name().unwrap_or("edsr")
+                );
+                std::process::exit(0);
+            }
+            _ => return Err(arg.unexpected()),
+        }
+    }
+
+    let model = values.pop_front().ok_or("missing `model` arg")?;
+    let image = values.pop_front().ok_or("missing `image` arg")?;
+    let output = values.pop_front().unwrap_or("output.png".into());
+
+    let args = Args {
+        image,
+        model,
+        output,
+    };
+
+    Ok(args)
+}
+
+/// Upscale an image using [EDSR](https://github.com/eugenesiow/super-image).
+///
+/// The ONNX model can be obtained by running the `export-edsr.py` script
+/// in this directory. You will need to install the `super-image` package
+/// first. See https://github.com/eugenesiow/super-image. After exporting
+/// the model convert it using:
+///
+/// ```
+/// rten-convert edsr.onnx
+/// cargo run --release --bin edsr edsr.rten image.jpg
+/// ```
+///
+/// This will generate `output.png`, a copy of the input image that has been
+/// upscaled by 2x.
+fn main() -> Result<(), Box<dyn Error>> {
+    let args = parse_args()?;
+    let model = Model::load_file(args.model)?;
+
+    let image: NdTensor<f32, 3> = read_image(&args.image)?;
+    let mut image = image.into_dyn();
+    image.insert_axis(0); // Add batch dim
+
+    let mut upscaled_image: Tensor<f32> = model.run_one(image.view().into(), None)?.try_into()?;
+    upscaled_image.remove_axis(0); // Remove batch dim
+
+    write_image(&args.output, upscaled_image.nd_view())?;
+
+    Ok(())
+}
diff --git a/rten-examples/src/export-edsr.py b/rten-examples/src/export-edsr.py
new file mode 100644
index 00000000..757512f8
--- /dev/null
+++ b/rten-examples/src/export-edsr.py
@@ -0,0 +1,28 @@
+# EDSR export script.
+#
+# Adapted from the example in the super-image README:
+# https://github.com/eugenesiow/super-image?tab=readme-ov-file#quick-start
+from PIL import Image
+import requests
+from super_image import EdsrModel, ImageLoader
+import torch
+
+url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
+
+model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)
+inputs = ImageLoader.load_image(image)
+
+torch.onnx.export(
+    # `model` is a `DataParallel` instance which can't be exported directly.
+    # Access the root module inside it.
+    model.module, 
+    (inputs,),
+    "edsr.onnx",
+    input_names=["input"],
+    output_names=["output"],
+    dynamic_axes={
+        "input": {0: "batch_size", 2: "height", 3: "width"},
+        "output": {0: "batch_size", 2: "height", 3: "width"},
+    },
+)