feat: refactor to proper app + write tests

src/BGR.Console/Removal/IImage.cs

@@ -4,6 +4,5 @@ internal interface IImage
 {
   int Width { get; }
   int Height { get; }
-  void Resize(int width, int height);
-  IPixel GetPixel(int x, int y);
+  Stream Data { get; }
 }

src/BGR.Console/Removal/IInferenceRunner.cs

@@ -0,0 +1,6 @@
+namespace BGR.Console.Removal;
+
+internal interface IInferenceRunner
+{
+  ITensor<float> Run(byte[] model, ITensor<float> inputTensor);
+}

src/BGR.Console/Removal/IPixel.cs

@@ -1,8 +0,0 @@
-namespace BGR.Console.Removal;
-
-internal interface IPixel
-{
-  float R { get; }
-  float G { get; }
-  float B { get; }
-}

src/BGR.Console/Removal/ITensor.cs

@@ -6,4 +6,5 @@ internal interface ITensor<T>
   int Width { get; }
   void SetValue(int batch, int channel, int y, int x, T value);
   float GetValue(int batch, int channel, int y, int x);
-}
+  Tensor<T> ToTensor();
+}

src/BGR.Console/Removal/ImageProcessor.cs

@@ -2,12 +2,16 @@ namespace BGR.Console.Removal;
 
 internal abstract class ImageProcessor
 {
+  public abstract Task<IImage> LoadImageAsync(string path);
+
   public abstract Task<ITensor<float>> CreateTensorInputAsync(Stream image, Model model);
 
-  public abstract Task<Stream> GenerateMaskAsync(OnnxTensor maskTensor, int width, int height);
+  public abstract Task<Stream> GenerateMaskAsync(ITensor<float> maskTensor, int width, int height);
 
   public abstract Task<Stream> RemoveBackgroundAsync(Stream image, Stream mask);
 
+  public abstract Task SaveImageAsync(Stream image, string path);
+
   protected static void WalkImage(int height, int width, Action<int, int> action)
   {
     for (var y = 0; y < height; y++)

src/BGR.Console/Removal/ImageSharp/ImageSharpProcessor.cs

@@ -2,9 +2,26 @@ namespace BGR.Console.Removal.ImageSharp;
 
 internal class ImageSharpProcessor : ImageProcessor
 {
+  public override async Task<IImage> LoadImageAsync(string path)
+  {
+    var image = await Image.LoadAsync<Rgba32>(path);
+
+    if (image.Metadata.DecodedImageFormat is null)
+    {
+      throw new InvalidOperationException("Image format is not supported.");
+    }
+
+    var stream = new MemoryStream();
+    await image.SaveAsync(stream, image.Metadata.DecodedImageFormat);
+    stream.Position = 0;
+
+    return new SharpImage(image.Width, image.Height, stream);
+  }
+
   public override async Task<ITensor<float>> CreateTensorInputAsync(Stream image, Model model)
   {
     using var resized = await Image.LoadAsync<Rgba32>(image);
+
     resized.Mutate(x => x.Resize(model.InputWidth, model.InputHeight));
 
     const int batchSize = 1;
@@ -22,7 +39,7 @@ internal class ImageSharpProcessor : ImageProcessor
     return tensor;
   }
 
-  public override async Task<Stream> GenerateMaskAsync(OnnxTensor maskTensor, int width, int height)
+  public override async Task<Stream> GenerateMaskAsync(ITensor<float> maskTensor, int width, int height)
   {
     using var mask = new Image<Rgba32>(width, height);
 
@@ -45,12 +62,16 @@ internal class ImageSharpProcessor : ImageProcessor
 
     var stream = new MemoryStream();
     await mask.SaveAsync(stream, new PngEncoder());
+    stream.Position = 0;
 
     return stream;
   }
 
   public override async Task<Stream> RemoveBackgroundAsync(Stream image, Stream mask)
   {
+    image.Position = 0;
+    mask.Position = 0;
+
     var imageWithBg = await Image.LoadAsync<Rgba32>(image);
     var maskImage = await Image.LoadAsync<Rgba32>(mask);
     using var imageWithBgRemoved = new Image<Rgba32>(imageWithBg.Width, imageWithBg.Height);
@@ -72,9 +93,18 @@ internal class ImageSharpProcessor : ImageProcessor
 
     var result = new MemoryStream();
     await imageWithBgRemoved.SaveAsync(result, new PngEncoder());
+    result.Position = 0;
+
     return result;
   }
 
+  public override async Task SaveImageAsync(Stream image, string path)
+  {
+    image.Position = 0;
+    var img = await Image.LoadAsync<Rgba32>(image);
+    await img.SaveAsync(path, new PngEncoder());
+  }
+
   private static float Normalize(float value)
   {
     const float binarizationThreshold = 0.5f;

src/BGR.Console/Removal/ImageSharp/SharpImage.cs

@@ -1,19 +1,25 @@
 namespace BGR.Console.Removal.ImageSharp;
 
-internal class SharpImage(Image<Rgba32> image) : IImage
+internal class SharpImage : IImage
 {
-  private readonly Image<Rgba32> _image = image;
+  public int Width { get; }
+  public int Height { get; }
+  public Stream Data { get; }
 
-  public int Width => _image.Width;
-  public int Height => _image.Height;
-
-  public void Resize(int width, int height)
+  public SharpImage(int width, int height, Stream data)
   {
-    _image.Mutate(x => x.Resize(width, height));
-  }
+    if (width <= 0)
+    {
+      throw new ArgumentOutOfRangeException(nameof(width), "must be greater than 0");
+    }
 
-  public IPixel GetPixel(int x, int y)
-  {
-    return new SharpPixel(_image[x, y]);
+    if (height <= 0)
+    {
+      throw new ArgumentOutOfRangeException(nameof(height), "must be greater than 0");
+    }
+
+    Width = width;
+    Height = height;
+    Data = data ?? throw new ArgumentNullException(nameof(data));
   }
 }

src/BGR.Console/Removal/ImageSharp/SharpPixel.cs

@@ -1,10 +0,0 @@
-namespace BGR.Console.Removal.ImageSharp;
-
-internal class SharpPixel(Rgba32 pixel) : IPixel
-{
-  private readonly Rgba32 _pixel = pixel;
-
-  public float R => _pixel.R;
-  public float G => _pixel.G;
-  public float B => _pixel.B;
-}

src/BGR.Console/Removal/Models/IModelFactory.cs

@@ -0,0 +1,6 @@
+namespace BGR.Console.Removal.Models;
+
+internal interface IModelFactory
+{
+  Model Create(string resourceName);
+}

src/BGR.Console/Removal/Models/ModNetModel.cs

@@ -1,7 +1,8 @@
 namespace BGR.Console.Removal.Models;
 
-internal class ModNetModel : Model
+internal class ModNetModel(byte[] modelBytes) : Model(modelBytes)
 {
+  public const string Id = "modnet";
   public override int InputWidth => 512;
   public override int InputHeight => 512;
   public override float RedNormalizationMean => 0.485f;

src/BGR.Console/Removal/Models/Model.cs

@@ -11,18 +11,28 @@ internal abstract class Model
   public abstract float RedNormalizationStd { get; }
   public abstract float GreenNormalizationStd { get; }
   public abstract float BlueNormalizationStd { get; }
+  public byte[] Bytes { get; } = [];
 
-  public float NormalizeRed(float value)
+  internal Model()
+  {
+  }
+
+  protected Model(byte[] modelBytes)
+  {
+    Bytes = modelBytes;
+  }
+
+  public virtual float NormalizeRed(float value)
   {
     return Normalize(value, RedNormalizationMean, RedNormalizationStd);
   }
 
-  public float NormalizeGreen(float value)
+  public virtual float NormalizeGreen(float value)
   {
     return Normalize(value, GreenNormalizationMean, GreenNormalizationStd);
   }
 
-  public float NormalizeBlue(float value)
+  public virtual float NormalizeBlue(float value)
   {
     return Normalize(value, BlueNormalizationMean, BlueNormalizationStd);
   }

src/BGR.Console/Removal/Models/ModelFactory.cs

@@ -0,0 +1,21 @@
+namespace BGR.Console.Removal.Models;
+
+internal class ModelFactory(IResourceManager resourceManager) : IModelFactory
+{
+  private readonly IResourceManager _resourceManager = resourceManager;
+
+  public Model Create(string resourceName)
+  {
+    var resource = _resourceManager.GetResource(resourceName);
+    var model = new byte[resource.Length];
+    resource.ReadExactly(model);
+
+    return resourceName switch
+    {
+      $"{U2NetModel.Id}.onnx" => new U2NetModel(model),
+      $"{RmbgModel.Id}.onnx" => new RmbgModel(model),
+      $"{ModNetModel.Id}.onnx" => new ModNetModel(model),
+      _ => throw new ArgumentException($"Unknown model name: {resourceName}")
+    };
+  }
+}

src/BGR.Console/Removal/Models/RmbgModel.cs

@@ -1,8 +1,8 @@
-
 namespace BGR.Console.Removal.Models;
 
-internal class RmbgModel : Model
+internal class RmbgModel(byte[] modelBytes) : Model(modelBytes)
 {
+  public const string Id = "rmbg";
   public override int InputWidth => 1024;
   public override int InputHeight => 1024;
   public override float RedNormalizationMean => 0.485f;

src/BGR.Console/Removal/Models/U2NetModel.cs

@@ -1,7 +1,8 @@
 namespace BGR.Console.Removal.Models;
 
-internal class U2NetModel : Model
+internal class U2NetModel(byte[] modelBytes) : Model(modelBytes)
 {
+  public const string Id = "u2net";
   public override int InputWidth => 320;
   public override int InputHeight => 320;
   public override float RedNormalizationMean => 0.485f;
@@ -10,4 +11,4 @@ internal class U2NetModel : Model
   public override float RedNormalizationStd => 0.229f;
   public override float GreenNormalizationStd => 0.224f;
   public override float BlueNormalizationStd => 0.225f;
-}
+}

src/BGR.Console/Removal/Onnx/OnnxInferenceRunner.cs

@@ -0,0 +1,18 @@
+namespace BGR.Console.Removal.Onnx;
+
+internal class OnnxInferenceRunner : IInferenceRunner
+{
+  public ITensor<float> Run(byte[] model, ITensor<float> inputTensor)
+  {
+    using var options = new SessionOptions() { LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_ERROR };
+    using var session = new InferenceSession(model, options);
+    var inputs = new List<NamedOnnxValue>()
+    {
+      NamedOnnxValue.CreateFromTensor(session.InputNames[0], inputTensor.ToTensor()),
+    };
+
+    var results = session.Run(inputs);
+    var outputTensor = results[0].AsTensor<float>();
+    return new OnnxTensor(outputTensor);
+  }
+}

src/BGR.Console/Removal/Onnx/OnnxTensor.cs

@@ -1,19 +1,28 @@
 namespace BGR.Console.Removal.Onnx;
 
-public class OnnxTensor(
-  int batchSize,
-  int channels,
-  int height,
-  int width
-) : ITensor<float>
+public class OnnxTensor : ITensor<float>
 {
-  private readonly DenseTensor<float> _tensor =
-    new([batchSize, channels, height, width]);
+  private readonly Tensor<float> _tensor;
 
   public int Height => _tensor.Dimensions[2];
 
   public int Width => _tensor.Dimensions[3];
 
+  public OnnxTensor(
+    int batchSize,
+    int channels,
+    int height,
+    int width
+  )
+  {
+    _tensor = new DenseTensor<float>([batchSize, channels, height, width]);
+  }
+
+  public OnnxTensor(Tensor<float> tensor)
+  {
+    _tensor = tensor;
+  }
+
   public void SetValue(
       int batch,
       int channel,
@@ -34,4 +43,9 @@ public class OnnxTensor(
   {
     return _tensor[batch, channel, y, x];
   }
+
+  public Tensor<float> ToTensor()
+  {
+    return _tensor;
+  }
 }

src/BGR.Console/Removal/RemovalCommand.cs

@@ -0,0 +1,139 @@
+using System.Diagnostics;
+
+namespace BGR.Console.Removal;
+
+internal class RemovalCommand(
+  IModelFactory modelFactory,
+  ImageProcessor imageProcessor,
+  IInferenceRunner inferenceRunner,
+  IAnsiConsole console,
+  ILogger<RemovalCommand> logger
+) : AsyncCommand<RemovalCommand.Settings>
+{
+  private readonly IModelFactory _modelFactory = modelFactory;
+  private readonly ImageProcessor _imageProcessor = imageProcessor;
+  private readonly IInferenceRunner _inferenceRunner = inferenceRunner;
+  private readonly IAnsiConsole _console = console;
+  private readonly ILogger<RemovalCommand> _logger = logger;
+
+  public override async Task<int> ExecuteAsync(CommandContext context, Settings settings)
+  {
+    await _console.Status()
+      .Spinner(Spinner.Known.Dots)
+      .SpinnerStyle(Style.Parse("green"))
+      .StartAsync("Removing background...", async ctx =>
+      {
+        ctx.Status("Loading model...");
+        var model = _logger.TimeAndLogAction(
+          "Loading model",
+          () => _modelFactory.Create(settings.ResourceName)
+        );
+
+        ctx.Status("Loading image...");
+        var image = await _logger.TimeAndLogActionAsync(
+          "Loading image",
+          async () => await _imageProcessor.LoadImageAsync(settings.Image)
+        );
+
+        ctx.Status("Creating tensor input...");
+        var inputTensor = await _logger.TimeAndLogActionAsync(
+          "Creating tensor input",
+          async () => await _imageProcessor.CreateTensorInputAsync(image.Data, model)
+        );
+
+        ctx.Status("Running inference...");
+        var outputTensor = _logger.TimeAndLogAction(
+          "Running inference",
+          () => _inferenceRunner.Run(model.Bytes, inputTensor)
+        );
+
+        ctx.Status("Generating mask...");
+        var mask = await _logger.TimeAndLogActionAsync(
+          "Generating mask",
+          async () => await _imageProcessor.GenerateMaskAsync(outputTensor, image.Width, image.Height)
+        );
+
+        ctx.Status("Removing background...");
+        var output = await _logger.TimeAndLogActionAsync(
+          "Removing background",
+          async () => await _imageProcessor.RemoveBackgroundAsync(image.Data, mask)
+        );
+
+        if (settings.IncludeMask)
+        {
+          await _logger.TimeAndLogActionAsync(
+            "Saving mask",
+            async () => await _imageProcessor.SaveImageAsync(mask, settings.MaskPath)
+          );
+
+          _console.MarkupLine($"[bold]Mask saved to:[/] [blue]{settings.OutputPath}[/]");
+        }
+
+        await _logger.TimeAndLogActionAsync(
+          "Saving output",
+          async () => await _imageProcessor.SaveImageAsync(output, settings.OutputPath)
+        );
+
+        _console.MarkupLine($"[bold]Output saved to:[/] [green]{settings.OutputPath}[/]");
+      });
+
+    return 0;
+  }
+
+  internal class Settings : CommandSettings
+  {
+    private static readonly Dictionary<string, string> Models = new()
+    {
+      { RmbgModel.Id, "rmbg.onnx" },
+      { ModNetModel.Id, "modnet.onnx" },
+      { U2NetModel.Id, "u2net.onnx" },
+    };
+
+    [CommandArgument(0, "<image>")]
+    [Description("Path to the image file whose background you want to remove")]
+    public string Image { get; init; } = string.Empty;
+
+    [CommandOption("--model|-m")]
+    [Description("The model to use for background removal")]
+    public string Model { get; init; } = "rmbg";
+
+    [CommandOption("--include-mask|-i")]
+    [Description("Generate and output the mask used for background removal")]
+    public bool IncludeMask { get; init; } = false;
+
+    [CommandOption("--output|-o")]
+    [Description("Path to output image without background to. File extension will always be .png")]
+    public string Output { get; init; } = string.Empty;
+
+    public string ResourceName => Models[Model];
+
+    public string MaskPath => GetOutputPath("_mask");
+
+    public string OutputPath => GetOutputPath("_no_bg");
+
+    public override ValidationResult Validate()
+    {
+      if (File.Exists(Image) is false)
+      {
+        return ValidationResult.Error($"The image file '{Image}' does not exist.");
+      }
+
+      if (Models.ContainsKey(Model) is false)
+      {
+        return ValidationResult.Error($"The model '{Model}' is not supported.");
+      }
+
+      return ValidationResult.Success();
+    }
+
+    private string GetOutputPath(string modifier)
+    {
+      if (string.IsNullOrWhiteSpace(Output))
+      {
+        return Path.ChangeExtension(Image, null) + modifier + ".png";
+      }
+
+      return Path.ChangeExtension(Output, ".png");
+    }
+  }
+}