diff --git a/d2vs/ocr.py b/d2vs/ocr.py index 25526d5..2758ac8 100644 --- a/d2vs/ocr.py +++ b/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, diff --git a/test/base.py b/test/base.py index 2825ca5..d5b63ce 100644 --- a/test/base.py +++ b/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 diff --git a/test/test_data/item_definition/shako.png b/test/test_data/item_definition/shako.png new file mode 100644 index 0000000..24a03ed Binary files /dev/null and b/test/test_data/item_definition/shako.png differ diff --git a/test/test_data/item_definition/tal_rashas_lidless_eye.png b/test/test_data/item_definition/tal_rashas_lidless_eye.png new file mode 100644 index 0000000..fe23cbc Binary files /dev/null and b/test/test_data/item_definition/tal_rashas_lidless_eye.png differ diff --git a/test/test_item_definitions.py b/test/test_item_definitions.py new file mode 100644 index 0000000..3564c20 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/test_simple_scanning.py b/test/test_simple_scanning.py index e0b40e2..f7b4af0 100644 --- a/test/test_simple_scanning.py +++ b/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")