Add image upscaling with waifu2x-ncnn-vulkan, async load images in scroll image view mode

3 files changed, 275 insertions, 27 deletions

diff --git a/src/ImageUtils.cpp b/src/ImageUtils.cpp
new file mode 100644
index 0000000..d160450
--- /dev/null
+++ b/src/ImageUtils.cpp
@@ -0,0 +1,114 @@
+#include "../include/ImageUtils.hpp"
+#include <string.h>
+#include <arpa/inet.h>
+
+namespace QuickMedia {
+    static bool gif_get_size(unsigned char *data, size_t data_size, int *width, int *height) {
+        if(data_size >= 10 && memcmp(data, "GIF", 3) == 0) {
+            *width = ((int)data[7] << 8) | data[6];
+            *height = ((int)data[9] << 8) | data[8];
+            return true;
+        }
+        return false;
+    }
+
+    static bool png_get_size(unsigned char *data, size_t data_size, int *width, int *height) {
+        if(data_size >= 24 && memcmp(data, "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A", 8) == 0) {
+            memcpy(width, data + 16, sizeof(int));
+            memcpy(height, data + 16 + sizeof(int), sizeof(int));
+            *width = htonl(*width);
+            *height = htonl(*height);
+            return true;
+        }
+        return false;
+    }
+
+    static bool is_cpu_little_endian() {
+        unsigned short i;
+        memcpy(&i, "LE", sizeof(i));
+        return (i & 0xF) == 'E';
+    }
+
+    static unsigned short read_uint16(unsigned char *data, bool cpu_little_endian, bool file_little_endian) {
+        unsigned short result;
+        memcpy(&result, data, sizeof(result));
+        if(cpu_little_endian != file_little_endian)
+            result = __builtin_bswap16(result);
+        return result;
+    }
+
+    static unsigned int read_uint32(unsigned char *data, bool cpu_little_endian, bool file_little_endian) {
+        unsigned int result;
+        memcpy(&result, data, sizeof(result));
+        if(cpu_little_endian != file_little_endian)
+            result = __builtin_bswap32(result);
+        return result;
+    }
+#if 0
+    static bool tiff_get_size(unsigned char *data, size_t data_size, int *width, int *height) {
+        if(data_size < 8)
+            return false;
+
+        bool cpu_little_endian = is_cpu_little_endian();
+        bool file_little_endian = true;
+        if(memcmp(data, "II", 2) == 0)
+            file_little_endian = true;
+        else if(memcmp(data, "MM", 2) == 0)
+            file_little_endian = false;
+        else
+            return false;
+
+        unsigned short id = read_uint16(data, cpu_little_endian, file_little_endian);
+        if(id != 42)
+            return false;
+
+        unsigned int offset_to_ifd = read_uint32(data, cpu_little_endian, file_little_endian);
+        if(offset_to_ifd )
+    }
+#endif
+
+    static bool jpeg_get_size(unsigned char *data, size_t data_size, int *width, int *height) {
+        if(data_size < 11 || memcmp(data, "\xFF\xD8\xFF\xE0", 4) != 0)
+            return false;
+
+        unsigned short block_length = ((int)data[4] << 8) | data[5];
+        if(memcmp(data + 6, "JFIF\0", 5) != 0)
+            return false;
+
+        size_t index = 4;
+        while(index < data_size) {
+            index += block_length;
+            if(index + 1 >= data_size) return false;
+            if(data[index] != 0xFF) return false; // Check if start of block
+            if((data[index + 1] == 0xC0 || data[index + 1] == 0xC2) && index + 8 < data_size) { // Start of frame marker
+                *height = ((int)data[index + 5] << 8) | data[index + 6];
+                *width = ((int)data[index + 7] << 8) | data[index + 8];
+                return true;
+            } else {
+                if(index + 2 >= data_size) return false;
+                index += 2;
+                block_length = ((int)data[index] << 8) | data[index + 1];
+            }
+        }
+
+        return false;
+    }
+
+    bool image_get_resolution(const Path &path, int *width, int *height) {
+        FILE *file = fopen(path.data.c_str(), "rb");
+        if(!file)
+            return false;
+
+        unsigned char data[512];
+        size_t bytes_read = fread(data, 1, sizeof(data), file);
+        fclose(file);
+
+        if(png_get_size(data, bytes_read, width, height))
+            return true;
+        else if(gif_get_size(data, bytes_read, width, height))
+            return true;
+        else if(jpeg_get_size(data, bytes_read, width, height))
+            return true;
+        return false;
+    }
+}
\ No newline at end of file
diff --git a/src/ImageViewer.cpp b/src/ImageViewer.cpp
index fc31274..86e7138 100644
--- a/src/ImageViewer.cpp
+++ b/src/ImageViewer.cpp
@@ -34,12 +34,34 @@ namespace QuickMedia {
         has_size_vertical_cursor = size_vertical_cursor.loadFromSystem(sf::Cursor::SizeVertical);
     }
 
+    void ImageViewer::load_image_async(const Path &path, std::shared_ptr<ImageData> image_data, int page) {
+        image_data->image_status = ImageStatus::LOADING;
+        image_data->texture.setSmooth(true);
+        assert(!loading_image);
+        loading_image = true;
+        image_loader_thread = std::thread([this, image_data, path, page]() {
+            std::string image_data_str;
+            if(file_get_content(path, image_data_str) == 0) {
+                if(image_data->texture.loadFromMemory(image_data_str.data(), image_data_str.size())) {
+                    image_data->sprite.setTexture(image_data->texture, true);
+                    page_size[page].size = get_page_size(page);
+                    page_size[page].loaded = true;
+                    image_data->image_status = ImageStatus::LOADED;
+                } else {
+                    image_data->image_status = ImageStatus::FAILED_TO_LOAD;
+                }
+            }
+            loading_image = false;
+        });
+        image_loader_thread.detach();
+    }
+
     bool ImageViewer::render_page(sf::RenderWindow &window, int page, double offset_y) {
         if(page < 0 || page >= (int)image_data.size())
             return false;
 
         const sf::Vector2<double> image_size = get_page_size(page);
-        std::unique_ptr<ImageData> &page_image_data = image_data[page];
+        std::shared_ptr<ImageData> &page_image_data = image_data[page];
         sf::Vector2<double> render_pos(std::floor(window_size.x * 0.5 - image_size.x * 0.5), - image_size.y * 0.5 + scroll + offset_y);
         if(render_pos.y + image_size.y <= 0.0 || render_pos.y >= window_size.y) {
             if(page_image_data)
@@ -58,8 +80,26 @@ namespace QuickMedia {
         }
 
         if(page_image_data) {
-            if(page_image_data->failed_to_load_image) {
-                sf::Text error_message("Failed to load image for page " + std::to_string(1 + page), *font, 30);
+            if(page_image_data->image_status == ImageStatus::LOADED) {
+                page_image_data->sprite.setPosition(render_pos.x, render_pos.y);
+                window.draw(page_image_data->sprite);
+            } else {
+                std::string page_str = std::to_string(1 + page);
+                std::string msg;
+                if(page_image_data->image_status == ImageStatus::WAITING) {
+                    if(!loading_image) {
+                        Path image_path = chapter_cache_dir;
+                        image_path.join(page_str);
+                        load_image_async(image_path, page_image_data, page);
+                    }
+                    msg = "Loading image for page " + page_str;
+                } else if(page_image_data->image_status == ImageStatus::LOADING) {
+                    msg = "Loading image for page " + page_str;
+                } else if(page_image_data->image_status == ImageStatus::FAILED_TO_LOAD) {
+                    msg = "Failed to load image for page " + page_str;
+                }
+
+                sf::Text error_message(std::move(msg), *font, 30);
                 auto text_bounds = error_message.getLocalBounds();
                 error_message.setFillColor(sf::Color::Black);
                 sf::Vector2<double> render_pos_text(std::floor(window_size.x * 0.5 - text_bounds.width * 0.5), - text_bounds.height * 0.5 + scroll + offset_y);
@@ -74,9 +114,6 @@ namespace QuickMedia {
 
                 error_message.setPosition(render_pos_text.x, render_pos_text.y);
                 window.draw(error_message);
-            } else {
-                page_image_data->sprite.setPosition(render_pos.x, render_pos.y);
-                window.draw(page_image_data->sprite);
             }
         } else {
             std::string page_str = std::to_string(1 + page);
@@ -103,23 +140,17 @@ namespace QuickMedia {
             // TODO: Make image loading asynchronous
             if(get_file_type(image_path) == FileType::REGULAR) {
                 fprintf(stderr, "ImageViewer: Loaded page %d\n", 1 + page);
-                page_image_data = std::make_unique<ImageData>();
+
+                page_image_data = std::make_shared<ImageData>();
                 page_image_data->visible_on_screen = true;
+                
                 std::string image_data;
                 if(file_get_content(image_path, image_data) == 0) {
-                    if(page_image_data->texture.loadFromMemory(image_data.data(), image_data.size())) {
-                        page_image_data->texture.setSmooth(true);
-                        page_image_data->sprite.setTexture(page_image_data->texture, true);
-                        //image_texture.generateMipmap();
-                        page_image_data->failed_to_load_image = false;
-                        page_size[page].size = get_page_size(page);
-                        page_size[page].loaded = true;
-                    } else {
-                        page_image_data->failed_to_load_image = true;
-                    }
+                    page_image_data->image_status = ImageStatus::WAITING;
+                    page_image_data->texture.setSmooth(true);
                 } else {
                     show_notification("Manga", "Failed to load image for page " + page_str + ". Image filepath: " + image_path.data, Urgency::CRITICAL);
-                    page_image_data->failed_to_load_image = true;
+                    page_image_data->image_status = ImageStatus::FAILED_TO_LOAD;
                 }
             }
 
@@ -289,7 +320,7 @@ namespace QuickMedia {
         for(auto &page_data : image_data) {
             if(page_data && !page_data->visible_on_screen) {
                 fprintf(stderr, "ImageViewer: Unloaded page %d\n", 1 + i);
-                page_data.reset();
+                page_data = nullptr;
             }
             ++i;
         }
diff --git a/src/QuickMedia.cpp b/src/QuickMedia.cpp
index 75bc027..6db58c9 100644
--- a/src/QuickMedia.cpp
+++ b/src/QuickMedia.cpp
@@ -13,6 +13,7 @@
 #include "../include/GoogleCaptcha.hpp"
 #include "../include/Notification.hpp"
 #include "../include/ImageViewer.hpp"
+#include "../include/ImageUtils.hpp"
 #include <cppcodec/base64_rfc4648.hpp>
 
 #include <SFML/Graphics/RectangleShape.hpp>
@@ -68,6 +69,11 @@ static int get_monitor_max_hz(Display *display) {
     return 60;
 }
 
+static void get_screen_resolution(Display *display, int *width, int *height) {
+    *width = DefaultScreenOfDisplay(display)->width;
+    *height = DefaultScreenOfDisplay(display)->height;
+}
+
 static bool has_gl_ext(Display *disp, const char *ext) {
     const char *extensions = glXQueryExtensionsString(disp, DefaultScreen(disp));
     if(!extensions)
@@ -193,6 +199,14 @@ namespace QuickMedia {
     }
 
     Program::~Program() {
+        running = false;
+        if(upscale_images) {
+            {
+                std::unique_lock<std::mutex> lock(image_upscale_mutex);
+                image_upscale_cv.notify_one();
+            }
+            image_upscale_thead.join();
+        }
         delete related_media_body;
         delete body;
         delete current_plugin;
@@ -223,6 +237,7 @@ namespace QuickMedia {
         fprintf(stderr, "  plugin                      The plugin to use. Should be either 4chan, manganelo, mangatown, mangadex, pornhub, youtube or dmenu\n");
         fprintf(stderr, "  --tor                       Use tor. Disabled by default\n");
         fprintf(stderr, "  --use-system-mpv-config     Use system mpv config instead of no config. Disabled by default\n");
+        fprintf(stderr, "  --upscale-images            Upscale low-resolution manga pages using waifu2x-ncnn-vulkan. Disabled by default\n");
         fprintf(stderr, "  -p                          Change the placeholder text for dmenu\n");
         fprintf(stderr, "EXAMPLES:\n");
         fprintf(stderr, "QuickMedia manganelo\n");
@@ -292,6 +307,8 @@ namespace QuickMedia {
                 use_tor = true;
             } else if(strcmp(argv[i], "--use-system-mpv-config") == 0) {
                 use_system_mpv_config = true;
+            } else if(strcmp(argv[i], "--upscale-images") == 0) {
+                upscale_images = true;
             } else if(strcmp(argv[i], "-p") == 0) {
                 if(i < argc - 1) {
                     search_placeholder = argv[i + 1];
@@ -309,6 +326,51 @@ namespace QuickMedia {
             return -2;
         }
 
+        if(upscale_images) {
+            if(!current_plugin->is_manga()) {
+                fprintf(stderr, "Option --upscale-images is only valid for manganelo, mangatown and mangadex\n");
+                return -2;
+            }
+
+            if(!is_program_executable_by_name("waifu2x-ncnn-vulkan")) {
+                fprintf(stderr, "waifu2x-ncnn-vulkan needs to be installed (and accessible from PATH environment variable) when using the --upscale-images option\n");
+                return -2;
+            }
+
+            running = true;
+            image_upscale_thead = std::thread([this]{
+                CopyOp copy_op;
+                while(running) {
+                    {
+                        std::unique_lock<std::mutex> lock(image_upscale_mutex);
+                        while(images_to_upscale.empty() && running) image_upscale_cv.wait(lock);
+                        if(!running)
+                            break;
+                        copy_op = images_to_upscale.front();
+                        images_to_upscale.pop_front();
+                    }
+
+                    Path tmp_file = copy_op.source;
+                    tmp_file.append(".tmp.jpg");
+
+                    fprintf(stderr, "Upscaling %s\n", copy_op.source.data.c_str());
+                    const char *args[] = { "waifu2x-ncnn-vulkan", "-i", copy_op.source.data.c_str(), "-o", tmp_file.data.c_str(), nullptr };
+                    if(exec_program(args, nullptr, nullptr) != 0) {
+                        fprintf(stderr, "Warning: failed to upscale %s with waifu2x-ncnn-vulkan\n", copy_op.source.data.c_str());
+                        // No conversion, but we need the file to have the destination name to see that the operation completed (and read it)
+                        if(rename(copy_op.source.data.c_str(), copy_op.destination.data.c_str()) != 0)
+                            perror(tmp_file.data.c_str());
+                        continue;
+                    }
+
+                    if(rename(tmp_file.data.c_str(), copy_op.destination.data.c_str()) != 0)
+                        perror(tmp_file.data.c_str());
+                }
+            });
+        } else {
+            running = true;
+        }
+
         current_plugin->use_tor = use_tor;
         window.setTitle("QuickMedia - " + current_plugin->name);
 
@@ -1828,6 +1890,7 @@ namespace QuickMedia {
         }
     }
 
+    // TODO: Cancel download when navigating to another non-manga page
     void Program::download_chapter_images_if_needed(Manga *image_plugin) {
         if(downloading_chapter_url == images_url)
             return;
@@ -1910,10 +1973,35 @@ namespace QuickMedia {
                     return false;
                 }
 
-                if(rename(image_filepath_tmp.data.c_str(), image_filepath.data.c_str()) != 0) {
-                    perror("rename");
-                    show_notification("Storage", "Failed to save image to file: " + image_filepath_tmp.data, Urgency::CRITICAL);
-                    return false;
+                bool rename_immediately = true;
+                if(upscale_images) {
+                    int screen_width, screen_height;
+                    get_screen_resolution(disp, &screen_width, &screen_height);
+
+                    int image_width, image_height;
+                    if(image_get_resolution(image_filepath_tmp, &image_width, &image_height)) {
+                        if(image_height < screen_height * 0.75) {
+                            rename_immediately = false;
+                            CopyOp copy_op;
+                            copy_op.source = image_filepath_tmp;
+                            copy_op.destination = image_filepath;
+                            std::unique_lock<std::mutex> lock(image_upscale_mutex);
+                            images_to_upscale.push_back(std::move(copy_op));
+                            image_upscale_cv.notify_one();
+                        } else {
+                            fprintf(stderr, "Info: not upscaling %s because the file is already large on your monitor (screen height: %d, image height: %d)\n", image_filepath_tmp.data.c_str(), screen_height, image_height);
+                        }
+                    } else {
+                        fprintf(stderr, "Warning: failed to upscale %s because QuickMedia failed to recognize the resolution of the image\n", image_filepath_tmp.data.c_str());
+                    }
+                }
+
+                if(rename_immediately) {
+                    if(rename(image_filepath_tmp.data.c_str(), image_filepath.data.c_str()) != 0) {
+                        perror(image_filepath_tmp.data.c_str());
+                        show_notification("Storage", "Failed to save image to file: " + image_filepath.data, Urgency::CRITICAL);
+                        return false;
+                    }
                 }
 
                 return true;
@@ -1922,6 +2010,7 @@ namespace QuickMedia {
     }
 
     void Program::image_page() {
+        image_download_cancel = false;
         search_bar->onTextUpdateCallback = nullptr;
         search_bar->onTextSubmitCallback = nullptr;
 
@@ -2024,25 +2113,25 @@ namespace QuickMedia {
                     if(event.key.code == sf::Keyboard::Up) {
                         if(image_index > 0) {
                             --image_index;
-                            return;
+                            goto end_of_images_page;
                         } else if(image_index == 0 && body->get_selected_item() < (int)body->items.size() - 1) {
                             // TODO: Make this work if the list is sorted differently than from newest to oldest.
                             body->filter_search_fuzzy("");
                             body->select_next_item();
                             select_episode(body->items[body->get_selected_item()].get(), true);
                             image_index = 99999; // Start at the page that shows we are at the end of the chapter
-                            return;
+                            goto end_of_images_page;
                         }
                     } else if(event.key.code == sf::Keyboard::Down) {
                         if(image_index < num_images) {
                             ++image_index;
-                            return;
+                            goto end_of_images_page;
                         } else if(image_index == num_images && body->get_selected_item() > 0) {
                             // TODO: Make this work if the list is sorted differently than from newest to oldest.
                             body->filter_search_fuzzy("");
                             body->select_previous_item();
                             select_episode(body->items[body->get_selected_item()].get(), true);
-                            return;
+                            goto end_of_images_page;
                         }
                     } else if(event.key.code == sf::Keyboard::Escape) {
                         current_page = Page::EPISODE_LIST;
@@ -2112,9 +2201,17 @@ namespace QuickMedia {
 
             window.display();
         }
+
+        end_of_images_page:
+        if(current_page != Page::IMAGES && current_page != Page::IMAGES_CONTINUOUS) {
+            image_download_cancel = true;
+            std::unique_lock<std::mutex> lock(image_upscale_mutex);
+            images_to_upscale.clear();
+        }
     }
 
     void Program::image_continuous_page() {
+        image_download_cancel = false;
         search_bar->onTextUpdateCallback = nullptr;
         search_bar->onTextSubmitCallback = nullptr;
 
@@ -2182,6 +2279,12 @@ namespace QuickMedia {
                 }
             }
         }
+
+        if(current_page != Page::IMAGES && current_page != Page::IMAGES_CONTINUOUS) {
+            image_download_cancel = true;
+            std::unique_lock<std::mutex> lock(image_upscale_mutex);
+            images_to_upscale.clear();
+        }
     }
 
     void Program::content_list_page() {