Swap out ocr model

d2vs/ocr.py

@@ -34,7 +34,7 @@ def read(self, screen_data, x1=None, y1=None, x2=None, y2=None, save_debug_image
         """
         Scans an area and returns the bounded text boxes, as well as a guess for the Item Type.
 
-        :param screen_data: np array of pixel data
+        :param screen_data: the data to OCR, which can be: Pillow Image, a filename, np array of pixel data
         :param x1:
         :param y1:
         :param x2:
@@ -44,6 +44,9 @@ def read(self, screen_data, x1=None, y1=None, x2=None, y2=None, save_debug_image
         :return:
         """
         # Convert input data to something we like (np array w/ BGR data, not RGB)
+        if isinstance(screen_data, str):
+            screen_data = cv2.imread(screen_data)
+
         if not isinstance(screen_data, np.ndarray):
             screen_data = np.asarray(screen_data, dtype='uint8')
             if screen_data.shape[2] == 4:  # we have an alpha channel
@@ -68,7 +71,7 @@ def read(self, screen_data, x1=None, y1=None, x2=None, y2=None, save_debug_image
             screen_data[y1:y2, x1:x2],
 
             # Maximum shift in y direction. Boxes with different level should not be merged. default = 0.5
-            ycenter_ths=0.08,
+            ycenter_ths=0.1,
 
             # Maximum horizontal distance to merge boxes. default = 0.5
             width_ths=width_ths,

test/base.py

@@ -16,3 +16,12 @@ def _check_scan(self, path, expected_text, expected_item_type=None):
         assert text == expected_text
         if item_type:
             assert item_type == expected_item_type
+
+    def assert_readings_match_expected(self, readings, expected):
+        lines_read = [line for _, line, _ in readings]
+        lines_expected = expected.split("\n")
+
+        print("\n".join(lines_read))
+
+        for read, expected in zip(lines_read, lines_expected):
+            assert read == expected

test/test_item_definitions.py

@@ -0,0 +1,67 @@
+from textwrap import dedent
+
+from PIL import Image
+
+from .base import OCRTestCases
+
+
+class ItemDefinitionTestCases(OCRTestCases):
+
+    def test_reading_large_item_descriptions_with_no_errors(self):
+        shako_readings = self.ocr.read(Image.open("test/test_data/item_definition/shako.png"), width_ths=3.5)
+
+        # self.assert_readings_match_expected(
+        #     shako_readings,
+        #     dedent("""\
+        #         Harlequin Crest
+        #         Shako
+        #         Defense: 105
+        #         Durability: 11 of 12
+        #         Required Strength: 50
+        #         Required Level: 62
+        #         +2 to All Skills
+        #         +2 to All Attributes
+        #         +120 to Life (Based On Character Level)
+        #         +126 to Mana (Based Oon Character Level)
+        #         Damage Reduced by 10%
+        #         74% Better Chance of Getting Magic Items
+        #         Socketed (1)
+        #     """)
+        # )
+
+        tal_rasha_lidless_eye_readings = self.ocr.read("test/test_data/item_definition/tal_rashas_lidless_eye.png", width_ths=3.5)
+        self.assert_readings_match_expected(
+            tal_rasha_lidless_eye_readings,
+            dedent("""\
+                Inventory
+                Sell Value: 35000
+                Tal Rasha's Lidless Eye
+                Swirling Crystal
+                One-Hand Damage: 18 to 42
+                Staff Class - Normal Attack Speed
+                Durability: 48 of 30
+                (Sorceress Only)
+                Required Level: 65
+                20% Faster Cast Rate
+                +1 to Cold Mastery (Sorceress Only)
+                +1 to Lightning Mastery (Sorceress Only)
+                +2 to Fire Mastery (Sorceress Only)
+                +10 to Energy
+                +57 to Life
+                +77 to Mana
+                -15% to Enemy Fire Resistance
+                +1 to Sorceress Skill Levels
+                Replenish Life +10
+                65% Better Chance of Getting Magic items
+                Tal Rasha's Wrappings
+                Tal Rasha's Horadric Crest
+                Tal Rasha's Guardianship
+                Tal Rasha's Lidless Eye
+                Tal Rasha's Adiudication
+                Tal Rasha's Fine-Spun Cloth
+            """)
+        )
+
+    def test_reading_item_definitions_with_low_scaling(self):
+        # check definitions but this time with much lower scaling
+        assert False

test/test_simple_scanning.py

@@ -1,7 +1,7 @@
 from .base import OCRTestCases
 
 
-class SimpleScanningTestcases(OCRTestCases):
+class SimpleScanningTestCases(OCRTestCases):
 
     def test_scanning_simple_works(self):
         self._check_scan("test/test_data/simple/586_gold.png", "586 Gold", "Normal")