Update item printer prize reader to read quantity

SerialPrograms/Source/PokemonSV/Inference/ItemPrinter/PokemonSV_ItemPrinterPrizeReader.cpp

@@ -8,10 +8,13 @@
 #include "CommonFramework/Logging/Logger.h"
 #include "CommonFramework/ImageTypes/ImageRGB32.h"
 #include "CommonFramework/ImageTools/ImageFilter.h"
+#include "CommonFramework/ImageTools/ImageStats.h"
 #include "CommonFramework/OCR/OCR_RawOCR.h"
 #include "CommonFramework/OCR/OCR_SmallDictionaryMatcher.h"
+#include "CommonFramework/OCR/OCR_NumberReader.h"
 #include "CommonFramework/VideoPipeline/VideoOverlayScopes.h"
 #include "PokemonSV_ItemPrinterPrizeReader.h"
+#include <iostream>
 
 namespace PokemonAutomation{
 namespace NintendoSwitch{
@@ -59,16 +62,20 @@ ItemPrinterPrizeReader::ItemPrinterPrizeReader(Language language)
     for (size_t c = 0; c < 10; c++){
         m_boxes_normal[c] = ImageFloatBox(0.361, 0.175 + c * 0.0764444, 0.240, 0.060);
         m_boxes_bonus[c] = ImageFloatBox(0.361, 0.204 + c * 0.0764444, 0.240, 0.060);
+        m_boxes_normal_quantity[c] = ImageFloatBox(0.649, 0.175 + c * 0.0764444, 0.034000, 0.060);
+        m_boxes_bonus_quantity[c] = ImageFloatBox(0.649, 0.204 + c * 0.0764444, 0.034000, 0.060);
     }
 }
 
 void ItemPrinterPrizeReader::make_overlays(VideoOverlaySet& items) const{
     for (size_t c = 0; c < 10; c++){
         items.add(COLOR_GREEN, m_boxes_normal[c]);
         items.add(COLOR_BLUE, m_boxes_bonus[c]);
+        items.add(COLOR_GREEN, m_boxes_normal_quantity[c]);
+        items.add(COLOR_BLUE, m_boxes_bonus_quantity[c]);
     }
 }
-std::array<std::string, 10> ItemPrinterPrizeReader::read(
+std::array<std::string, 10> ItemPrinterPrizeReader::read_prizes(
     Logger& logger, AsyncDispatcher& dispatcher,
     const ImageViewRGB32& screen
 ) const{
@@ -117,7 +124,75 @@ std::array<std::string, 10> ItemPrinterPrizeReader::read(
 }
 
 
+std::array<int16_t, 10> ItemPrinterPrizeReader::read_quantity(
+    Logger& logger, AsyncDispatcher& dispatcher,
+    const ImageViewRGB32& screen
+) const{
+    size_t total_rows = 10;
+
+    // determine whether to use normal or bonus boxes for OCR
+    double total_average_sum_normal = 0;
+    double total_average_sum_bonus = 0;
+    for (size_t i = 0; i < total_rows; i++){
+        total_average_sum_normal += average_sum_filtered(screen, m_boxes_normal[i]);
+        total_average_sum_bonus += average_sum_filtered(screen, m_boxes_bonus[i]);
+    }
 
+    // std::cout << "total_average_sum_normal: " << std::to_string(total_average_sum_normal) << std::endl;
+    // std::cout << "total_average_sum_bonus: " << std::to_string(total_average_sum_bonus) << std::endl;
+
+    const std::array<ImageFloatBox, 10>& boxes = (total_average_sum_normal > total_average_sum_bonus) 
+                                                    ? m_boxes_bonus_quantity 
+                                                    : m_boxes_normal_quantity;
+
+    std::array<int16_t, 10> results;
+    std::vector<std::unique_ptr<AsyncTask>> tasks(total_rows);
+    for (size_t i = 0; i < total_rows; i++){
+        // ImageRGB32 filtered = to_blackwhite_rgb32_range(
+        //     extract_box_reference(screen, m_boxes_bonus_quantity[i]),
+        //     0xff808000, 0xffffffff,
+        //     true
+        // );
+        // filtered.save("DebugDumps/test"+ std::to_string(i) +".png");   
+
+        tasks[i] = dispatcher.dispatch([&, i]{
+            results[i] = read_number(logger, screen, boxes[i]);
+        });
+    }
+
+    for (size_t c = 0; c < total_rows; c++){
+        tasks[c]->wait_and_rethrow_exceptions();
+    }
+
+    return results;
+}
+
+
+int16_t ItemPrinterPrizeReader::read_number(
+    Logger& logger, 
+    const ImageViewRGB32& screen, 
+    const ImageFloatBox& box
+) const{
+
+    ImageViewRGB32 cropped = extract_box_reference(screen, box);
+    int16_t number = (int16_t)OCR::read_number_waterfill(logger, cropped, 0xff808000, 0xffffffff);
+
+    if (number < 1 || number > 40){
+        number = 1; // default to 1 if we can't read the prize quantity
+    }
+    return (int16_t)number;
+}
+
+double ItemPrinterPrizeReader::average_sum_filtered(const ImageViewRGB32& screen, const ImageFloatBox& box) const{
+    ImageViewRGB32 cropped = extract_box_reference(screen, box);
+    ImageRGB32 filtered = to_blackwhite_rgb32_range(
+        cropped,
+        0xff808000, 0xffffffff,
+        false
+    );    
+
+    return image_stats(filtered).average.sum();
+}
 
 
 }

SerialPrograms/Source/PokemonSV/Inference/ItemPrinter/PokemonSV_ItemPrinterPrizeReader.h

@@ -29,15 +29,31 @@ class ItemPrinterPrizeReader{
     ItemPrinterPrizeReader(Language language);
 
     void make_overlays(VideoOverlaySet& items) const;
-    std::array<std::string, 10> read(
+
+    std::array<std::string, 10> read_prizes(
+        Logger& logger, AsyncDispatcher& dispatcher,
+        const ImageViewRGB32& screen
+    ) const;
+
+    std::array<int16_t, 10> read_quantity(
         Logger& logger, AsyncDispatcher& dispatcher,
         const ImageViewRGB32& screen
     ) const;
 
+    int16_t read_number(
+        Logger& logger,
+        const ImageViewRGB32& screen, 
+        const ImageFloatBox& box
+    ) const;
+
+    double average_sum_filtered(const ImageViewRGB32& screen, const ImageFloatBox& box) const;
+
 private:
     Language m_language;
-    ImageFloatBox m_boxes_normal[10];
-    ImageFloatBox m_boxes_bonus[10];
+    std::array<ImageFloatBox, 10> m_boxes_normal;
+    std::array<ImageFloatBox, 10> m_boxes_bonus;
+    std::array<ImageFloatBox, 10> m_boxes_normal_quantity;
+    std::array<ImageFloatBox, 10> m_boxes_bonus_quantity;
 };
 
 

SerialPrograms/Source/PokemonSV/Programs/ItemPrinter/PokemonSV_ItemPrinterTools.cpp

@@ -116,7 +116,7 @@ std::array<std::string, 10> item_printer_finish_print(
                 VideoOverlaySet overlays(console.overlay());
                 reader.make_overlays(overlays);
                 auto snapshot = console.video().snapshot();
-                ret = reader.read(console.logger(), dispatcher, snapshot);
+                ret = reader.read_prizes(console.logger(), dispatcher, snapshot);
 //                static int c = 0;
 //                snapshot->save("test-" + std::to_string(c) + ".png");
             }