diff --git a/src/cli/cli.cpp b/src/cli/cli.cpp
index 7db6cc5..323db76 100644
--- a/src/cli/cli.cpp
+++ b/src/cli/cli.cpp
@@ -4,11 +4,11 @@
 
 #include <fmt/chrono.h>
 #include <fmt/format.h>
-#include <stb_image_write.h>
 
-#include "lib/assets/scene_parser.hpp"
-#include "lib/camera.hpp"
-#include "lib/path_tracer.hpp"
+#include "../lib/assets/scene_parser.hpp"
+#include "../lib/camera.hpp"
+#include "../lib/image.hpp"
+#include "../lib/path_tracer.hpp"
 
 [[nodiscard]] auto now() -> std::chrono::steady_clock::time_point
 {
@@ -58,7 +58,7 @@ class Stopwatch {
 };
 
 void execute_cli_version(const CliConfigurations& configs,
-                         std::filesystem::path asset_path)
+                         const std::filesystem::path& asset_path)
 {
   Stopwatch stopwatch;
 
@@ -103,12 +103,10 @@ void execute_cli_version(const CliConfigurations& configs,
   path_tracer.send_to_preview(buffer.data(), resolution);
 
   CUDA_CHECK(cudaDeviceSynchronize());
-  stopwatch.end_stage("Post Processing");
 
-  if (stbi_write_png(configs.output_filename.value().c_str(), width, height, 4,
-                     buffer.data(), 0) == 0) {
-    fmt::print(stderr, "Filed to write to file output.png");
-  }
+  write_image_file(configs.output_filename.value(), resolution.to_signed(),
+                   buffer.data());
+
   stopwatch.end_stage("Write image file");
 
   fmt::print("Done path tracing {}!\n\n", scene_desc.filename);
diff --git a/src/cli/cli.hpp b/src/cli/cli.hpp
index 29b1364..f5429e3 100644
--- a/src/cli/cli.hpp
+++ b/src/cli/cli.hpp
@@ -6,6 +6,6 @@
 #include <filesystem>
 
 void execute_cli_version(const CliConfigurations& configs,
-                         std::filesystem::path asset_path);
+                         const std::filesystem::path& asset_path);
 
 #endif // CUDA_PATH_TRACER_CLI_HPP
diff --git a/src/lib/CMakeLists.txt b/src/lib/CMakeLists.txt
index ae565e7..dff8bb3 100644
--- a/src/lib/CMakeLists.txt
+++ b/src/lib/CMakeLists.txt
@@ -49,7 +49,7 @@ add_library(cuda_pt_lib
         cuda_utils/definitions.hpp
         transform.hpp
         configurations.cpp
-        configurations.hpp resolution.hpp third_party_impl.cpp prelude.cpp prelude.hpp assets/assets.cpp assets/assets.hpp hash.cuh aabb.hpp assets/model_loader.cpp assets/model_loader.hpp static_stack.hpp)
+        configurations.hpp resolution.hpp third_party_impl.cpp prelude.cpp prelude.hpp assets/assets.cpp assets/assets.hpp hash.cuh aabb.hpp assets/model_loader.cpp assets/model_loader.hpp static_stack.hpp image.cpp image.hpp)
 add_library(cuda_pt::lib ALIAS cuda_pt_lib)
 
 target_link_libraries(cuda_pt_lib PUBLIC
diff --git a/src/lib/image.cpp b/src/lib/image.cpp
new file mode 100644
index 0000000..f0c30d4
--- /dev/null
+++ b/src/lib/image.cpp
@@ -0,0 +1,20 @@
+#include <filesystem>
+
+#include "image.hpp"
+
+#include <stb_image_write.h>
+
+void write_image_file(std::string filename, Resolution res, const void* data)
+{
+  auto [width, height] = res;
+
+  std::filesystem::path output_path{filename};
+
+  if (output_path.extension() == ".png") {
+    if (stbi_write_png(filename.c_str(), width, height, 4, data, 0) == 0) {
+      SPDLOG_CRITICAL("Failed to write to image file {}", filename);
+    }
+  } else {
+    fmt::print(stderr, "{} has an unrecognized extension\n", filename);
+  }
+}
\ No newline at end of file
diff --git a/src/lib/image.hpp b/src/lib/image.hpp
new file mode 100644
index 0000000..e846de8
--- /dev/null
+++ b/src/lib/image.hpp
@@ -0,0 +1,10 @@
+#ifndef CUDA_PATH_TRACER_IMAGE_HPP
+#define CUDA_PATH_TRACER_IMAGE_HPP
+
+#include <string>
+
+#include "resolution.hpp"
+
+void write_image_file(std::string filename, Resolution res, const void* data);
+
+#endif // CUDA_PATH_TRACER_IMAGE_HPP