wip image edit api

examples/common/common.hpp

@@ -1,4 +1,5 @@
 
+#include <filesystem>
 #include <iostream>
 #include <map>
 #include <random>
@@ -6,10 +7,9 @@
 #include <sstream>
 #include <string>
 #include <vector>
-#include <filesystem>
 
 #include <json.hpp>
-using json = nlohmann::json;
+using json   = nlohmann::json;
 namespace fs = std::filesystem;
 
 #if defined(_WIN32)
@@ -19,6 +19,18 @@ namespace fs = std::filesystem;
 
 #include "stable-diffusion.h"
 
+#define STB_IMAGE_IMPLEMENTATION
+#define STB_IMAGE_STATIC
+#include "stb_image.h"
+
+#define STB_IMAGE_WRITE_IMPLEMENTATION
+#define STB_IMAGE_WRITE_STATIC
+#include "stb_image_write.h"
+
+#define STB_IMAGE_RESIZE_IMPLEMENTATION
+#define STB_IMAGE_RESIZE_STATIC
+#include "stb_image_resize.h"
+
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
@@ -1612,3 +1624,123 @@ struct SDGenerationParams {
 static std::string version_string() {
     return std::string("stable-diffusion.cpp version ") + sd_version() + ", commit " + sd_commit();
 }
+
+uint8_t* load_image_common(bool from_memory,
+                           const char* image_path_or_bytes,
+                           int& width,
+                           int& height,
+                           int expected_width   = 0,
+                           int expected_height  = 0,
+                           int expected_channel = 3) {
+    int c = 0;
+    const char* image_path;
+    uint8_t* image_buffer = nullptr;
+    if (from_memory) {
+        image_path   = "memory";
+        image_buffer = (uint8_t*)stbi_load(image_path_or_bytes, &width, &height, &c, expected_channel);
+    } else {
+        image_path   = image_path_or_bytes;
+        image_buffer = (uint8_t*)stbi_load(image_path_or_bytes, &width, &height, &c, expected_channel);
+    }
+    if (image_buffer == nullptr) {
+        fprintf(stderr, "load image from '%s' failed\n", image_path);
+        return nullptr;
+    }
+    if (c < expected_channel) {
+        fprintf(stderr,
+                "the number of channels for the input image must be >= %d,"
+                "but got %d channels, image_path = %s\n",
+                expected_channel,
+                c,
+                image_path);
+        free(image_buffer);
+        return nullptr;
+    }
+    if (width <= 0) {
+        fprintf(stderr, "error: the width of image must be greater than 0, image_path = %s\n", image_path);
+        free(image_buffer);
+        return nullptr;
+    }
+    if (height <= 0) {
+        fprintf(stderr, "error: the height of image must be greater than 0, image_path = %s\n", image_path);
+        free(image_buffer);
+        return nullptr;
+    }
+
+    // Resize input image ...
+    if ((expected_width > 0 && expected_height > 0) && (height != expected_height || width != expected_width)) {
+        float dst_aspect = (float)expected_width / (float)expected_height;
+        float src_aspect = (float)width / (float)height;
+
+        int crop_x = 0, crop_y = 0;
+        int crop_w = width, crop_h = height;
+
+        if (src_aspect > dst_aspect) {
+            crop_w = (int)(height * dst_aspect);
+            crop_x = (width - crop_w) / 2;
+        } else if (src_aspect < dst_aspect) {
+            crop_h = (int)(width / dst_aspect);
+            crop_y = (height - crop_h) / 2;
+        }
+
+        if (crop_x != 0 || crop_y != 0) {
+            printf("crop input image from %dx%d to %dx%d, image_path = %s\n", width, height, crop_w, crop_h, image_path);
+            uint8_t* cropped_image_buffer = (uint8_t*)malloc(crop_w * crop_h * expected_channel);
+            if (cropped_image_buffer == nullptr) {
+                fprintf(stderr, "error: allocate memory for crop\n");
+                free(image_buffer);
+                return nullptr;
+            }
+            for (int row = 0; row < crop_h; row++) {
+                uint8_t* src = image_buffer + ((crop_y + row) * width + crop_x) * expected_channel;
+                uint8_t* dst = cropped_image_buffer + (row * crop_w) * expected_channel;
+                memcpy(dst, src, crop_w * expected_channel);
+            }
+
+            width  = crop_w;
+            height = crop_h;
+            free(image_buffer);
+            image_buffer = cropped_image_buffer;
+        }
+
+        printf("resize input image from %dx%d to %dx%d\n", width, height, expected_width, expected_height);
+        int resized_height = expected_height;
+        int resized_width  = expected_width;
+
+        uint8_t* resized_image_buffer = (uint8_t*)malloc(resized_height * resized_width * expected_channel);
+        if (resized_image_buffer == nullptr) {
+            fprintf(stderr, "error: allocate memory for resize input image\n");
+            free(image_buffer);
+            return nullptr;
+        }
+        stbir_resize(image_buffer, width, height, 0,
+                     resized_image_buffer, resized_width, resized_height, 0, STBIR_TYPE_UINT8,
+                     expected_channel, STBIR_ALPHA_CHANNEL_NONE, 0,
+                     STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
+                     STBIR_FILTER_BOX, STBIR_FILTER_BOX,
+                     STBIR_COLORSPACE_SRGB, nullptr);
+        width  = resized_width;
+        height = resized_height;
+        free(image_buffer);
+        image_buffer = resized_image_buffer;
+    }
+    return image_buffer;
+}
+
+uint8_t* load_image_from_file(const char* image_path,
+                              int& width,
+                              int& height,
+                              int expected_width   = 0,
+                              int expected_height  = 0,
+                              int expected_channel = 3) {
+    return load_image_common(false, image_path, width, height, expected_width, expected_height, expected_channel);
+}
+
+uint8_t* load_image_from_memory(const char* image_bytes,
+                                int& width,
+                                int& height,
+                                int expected_width   = 0,
+                                int expected_height  = 0,
+                                int expected_channel = 3) {
+    return load_image_common(true, image_bytes, width, height, expected_width, expected_height, expected_channel);
+}