├── .gitignore ├── LICENSE.txt ├── README.md ├── d4-item-tooltip-ocr.py ├── examples ├── screenshot_001.png └── tooltip_001.jpg ├── paddleocr-models └── en_PP-OCRv3_rec-d4_tooltip │ ├── inference.pdiparams │ ├── inference.pdiparams.info │ └── inference.pdmodel ├── requirements.txt └── templates ├── affix.png ├── enchanted_rerolled.png ├── inprint_aspect.png ├── socket.png ├── socket_mask.png ├── socket_mask_new.png └── weapon_stat.png /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ 161 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 ARK Mod 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Diablo IV: Item Tooltip OCR 2 | 3 | Utilizing PaddleOCR and a custom Diablo 4 trained recognition model. Outputs item tooltip data in JSON-format. 4 | 5 | 6 | ## Example 7 | 8 | ![item tooltip](https://github.com/mxtsdev/d4-item-tooltip-ocr/assets/58796811/06aa5a7d-23a4-40c3-9d57-1916d3d3d32b) 9 | 10 | ```json 11 | { 12 | "affixes": [ 13 | "+22.5% Overpower Damage [22.5]%", 14 | "+12.5% Damage to Slowed Enemies [11.5 - 16.5]%", 15 | "+14.0% Critical Strike Damage [10.0 - 15.0]%", 16 | "+44 Willpower +[41 - 51]", 17 | "+7.0% Damage Over Time [5.0 - 10.0]%" 18 | ], 19 | "aspect": "Core Skills deal an additional 7.0%[x] [6.0 - 8.0]% damage for each active Companion. (Druid Only)", 20 | "item_power": "710", 21 | "item_power_upgraded": null, 22 | "item_upgrades_current": null, 23 | "item_upgrades_max": null, 24 | "name": "SHEPHERD'S WOLF'S BITE", 25 | "sockets": [], 26 | "stats": [ 27 | "806 Damage Per Second (-1555)", 28 | "[586 - 880] Damage per Hit", 29 | "1.10 Attacks per Second (Fast Weapon)" 30 | ], 31 | "type": "Sacred Legendary Mace" 32 | } 33 | ``` 34 | 35 | 36 | ## Installation 37 | 38 | - Clone repository 39 | 40 | - Create a Python3 enviroment (I recommend using https://www.anaconda.com/download on Windows) 41 | 42 | - pip install -r requirements.txt 43 | 44 | You will need to install the correct version of PaddlePaddle depending on your environment (CPU/GPU/CUDA version). Please refer to this link: 45 | https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/en/install/pip/windows-pip_en.html#old-version-anchor-3-INSTALLATION 46 | 47 | ## Using 48 | 49 | ### Output json to console 50 | ``` 51 | python d4-item-tooltip-ocr.py --source-img=examples\screenshot_001.png 52 | python d4-item-tooltip-ocr.py --source-img=examples\tooltip_001.jpg --find-tooltip=False 53 | ``` 54 | 55 | ### Output json to file 56 | ``` 57 | python d4-item-tooltip-ocr.py --source-img=examples\screenshot_001.png --json-output=item-tooltip-data.json 58 | python d4-item-tooltip-ocr.py --source-img=examples\tooltip_001.jpg --json-output=item-tooltip-data.json --find-tooltip=False 59 | ``` 60 | 61 | ### Debug mode 62 | ``` 63 | python d4-item-tooltip-ocr.py --debug=True --source-img=examples\screenshot_001.png 64 | python d4-item-tooltip-ocr.py --debug=True --source-img=examples\tooltip_001.jpg --find-tooltip=False 65 | ``` 66 | -------------------------------------------------------------------------------- /d4-item-tooltip-ocr.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import json 4 | import re 5 | import time 6 | import os 7 | import tkinter as tk 8 | import argparse 9 | from paddleocr import PaddleOCR, draw_ocr 10 | from Levenshtein import ratio 11 | 12 | class LineEntry(): 13 | def __init__(self, index, txt, boxes, score): 14 | self.index = index 15 | self.txt = txt 16 | self.score = score 17 | x = [b[0] for b in boxes] 18 | y = [b[1] for b in boxes] 19 | self.tl = (min(x), min(y)) #top-left 20 | self.br = (max(x), max(y)) #bottom-right 21 | self.dx = self.br[0] - self.tl[0] 22 | self.dy = self.br[1] - self.tl[1] 23 | self.cy = self.tl[1] + int(self.dy / 2.0) 24 | self.a = self.dx * self.dy 25 | 26 | @staticmethod 27 | def EMPTY(): 28 | return LineEntry(-1, None, [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], 0.0) 29 | 30 | def re_get(self, pattern: str, *return_group_names, flags = 0): 31 | if self.index == -1: return (*[None for gn in return_group_names],) 32 | 33 | m = re.search(pattern, self.txt) 34 | results = [m.group(gn) for gn in return_group_names] 35 | return (*results,) 36 | 37 | class LineCollection(): 38 | def __init__(self, result): 39 | if isinstance(result, list) and all(isinstance(x, LineEntry) for x in result): 40 | self.lines = result 41 | elif isinstance(result, list) and len(result) == 0: 42 | self.lines = [] 43 | else: 44 | self.lines = list(filter(lambda x : x.a > 150, [LineEntry(i, line[1][0], line[0], line[1][1]) for i, line in enumerate(result)])) 45 | 46 | def __len__(self): 47 | return len(self.lines) 48 | 49 | def __getitem__(self, key): 50 | return self.lines[key] 51 | 52 | # find string index using basic string compare and fallback to levenshtein distance 53 | def find(self, key, score_cutoff=0.5, strip=None, strip_flags = 0): 54 | # basic string compare 55 | for i, line in enumerate(self.lines): 56 | index = line.txt.find(key) 57 | if index != -1: 58 | return i 59 | 60 | # compare using levenshtein distance 61 | for i, line in enumerate(self.lines): 62 | str = line.txt 63 | if strip is not None: 64 | str = re.sub(strip, '', str, flags=strip_flags) 65 | 66 | r = ratio(str, key, score_cutoff=score_cutoff) 67 | if r >= score_cutoff: 68 | return i 69 | 70 | return -1 71 | 72 | # take lines while matching predicate, optionally remove from source 73 | def takewhile(self, predicate, remove=True): 74 | lines = [] 75 | for i in self.lines: 76 | if predicate(i): 77 | lines.append(i) 78 | else: 79 | break 80 | 81 | if remove: 82 | self.lines = self.lines[len(lines):] 83 | 84 | return LineCollection(lines) 85 | 86 | # take lines between given cy offsets 87 | def takebetween(self, y_min, y_max, remove=True): 88 | lines = [] 89 | for i in self.lines: 90 | if y_min <= i.cy and y_max >= i.cy: 91 | lines.append(i) 92 | 93 | if remove: 94 | for i in lines: 95 | self.lines.remove(i) 96 | 97 | return LineCollection(lines) 98 | 99 | # split line collection at offset (returning that entry, collection of preceding lines, collection of trailing lines) 100 | def splitat(self, index): 101 | if index < 0 or index > len(self.lines) - 1: 102 | return LineEntry.EMPTY(), LineCollection([]), LineCollection([]) 103 | 104 | return self.lines[index], LineCollection(self.lines[0:index]), LineCollection(self.lines[index+1:]) 105 | 106 | # used in debug mode to navigate between images in a directory 107 | class SimpleDirectoryNavigator(): 108 | def __init__(self, source_path, run_tests_mode=False): 109 | self.run_tests_mode = run_tests_mode 110 | 111 | if run_tests_mode == False: 112 | self.dir_path = os.path.dirname(source_path) 113 | self.current_image_filename = os.path.basename(source_path) 114 | else: 115 | self.dir_path = source_path 116 | self.current_image_filename = None 117 | 118 | def getImagePath(self, offsetFromCurrent = 0): 119 | for root, dirs, files in os.walk(self.dir_path): 120 | imgs = list(filter(lambda x : x.endswith('.png') or x.endswith('.jpg'), files)) 121 | for i, file in enumerate(imgs): 122 | if file == self.current_image_filename: 123 | index = i + offsetFromCurrent 124 | if index >= 0 and index < len(imgs): 125 | self.current_image_filename = imgs[index] 126 | return os.path.join(self.dir_path, self.current_image_filename), \ 127 | index - 1 >= 0 and index - 1 < len(imgs), \ 128 | index + 1 >= 0 and index + 1 < len(imgs) 129 | else: 130 | return None, \ 131 | index - 1 >= 0 and index - 1 < len(imgs), \ 132 | index + 1 >= 0 and index + 1 < len(imgs) 133 | break 134 | 135 | def getImagePath_RunTestsMode(self, offsetFromCurrent = 0): 136 | def isTestFile(x): 137 | if not (x.endswith('.png') or x.endswith('.jpg')): 138 | return False 139 | 140 | filename_without_ext = os.path.splitext(x)[0] 141 | json_path = os.path.join(self.dir_path, filename_without_ext + '.json') 142 | 143 | return os.path.isfile(json_path) 144 | 145 | for root, dirs, files in os.walk(self.dir_path): 146 | imgs = list(filter(isTestFile, files)) 147 | for i, file in enumerate(imgs): 148 | if self.current_image_filename == None: 149 | index = i 150 | self.current_image_filename = imgs[index] 151 | return os.path.join(self.dir_path, self.current_image_filename), \ 152 | index - 1 >= 0 and index - 1 < len(imgs), \ 153 | index + 1 >= 0 and index + 1 < len(imgs) 154 | 155 | if file == self.current_image_filename: 156 | index = i + offsetFromCurrent 157 | if index >= 0 and index < len(imgs): 158 | self.current_image_filename = imgs[index] 159 | return os.path.join(self.dir_path, self.current_image_filename), \ 160 | index - 1 >= 0 and index - 1 < len(imgs), \ 161 | index + 1 >= 0 and index + 1 < len(imgs) 162 | else: 163 | return None 164 | break 165 | 166 | class D4ItemTooltipOCR(): 167 | def __init__(self): 168 | self.ocr = PaddleOCR( 169 | lang='en', 170 | use_angle_cls=False, 171 | show_log=False, 172 | det_db_unclip_ratio=2.0, 173 | rec_model_dir='paddleocr-models/en_PP-OCRv3_rec-d4_tooltip', 174 | rec_batch_num=10, 175 | enable_mkldnn=True) 176 | self.img_tmpl_affix = cv2.imread('templates/affix.png') 177 | self.img_tmpl_reroll = cv2.imread('templates/enchanted_rerolled.png') 178 | self.img_tmpl_aspect = cv2.imread('templates/inprint_aspect.png') 179 | self.img_tmpl_wstat = cv2.imread('templates/weapon_stat.png') 180 | self.img_tmpl_socket = cv2.imread('templates/socket.png') 181 | self.img_tmpl_socket_mask = cv2.imread('templates/socket_mask_new.png') 182 | 183 | self.templates = { 184 | 'affix': self.img_tmpl_affix, 185 | 'reroll': self.img_tmpl_reroll, 186 | 'aspect': self.img_tmpl_aspect, 187 | 'wstat': self.img_tmpl_wstat, 188 | 'socket': [self.img_tmpl_socket, self.img_tmpl_socket_mask] 189 | } 190 | 191 | def processImage(self, source_path, find_tooltip=True, debug=False): 192 | start_pi = time.time() 193 | jsonstr, ocr_deltatime = self.processImage_internal(source_path, find_tooltip, debug) 194 | end_pi = time.time() 195 | print(f'Total: {end_pi-start_pi:.2f}s, OCR: {ocr_deltatime:.2f}s', end='\n\n') 196 | 197 | return jsonstr 198 | 199 | def processImage_internal(self, source_path, find_tooltip=True, debug=False): 200 | if debug == True: 201 | print(f'[Source image: \'{source_path}\']') 202 | 203 | input_image = cv2.imread(source_path) 204 | 205 | # find the item tooltip 206 | found_tooltip = False 207 | if find_tooltip == True: 208 | # preprocess input image 209 | height, width = input_image.shape[:2] 210 | hsv = cv2.cvtColor(input_image, cv2.COLOR_BGR2HSV) 211 | tmp = cv2.inRange(hsv, np.array([69, 45, 47]), np.array([85, 106, 73])) 212 | 213 | kernel = np.ones((8,8), np.uint8) 214 | tmp = cv2.dilate(tmp, kernel, iterations = 5) 215 | 216 | # find contour most likely to be the tooltip 217 | tooltipcontours_image = input_image.copy() 218 | contours, hierarchy = cv2.findContours(tmp, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) 219 | for cnt in contours: 220 | x, y, w, h = cv2.boundingRect(cnt) 221 | if cv2.contourArea(cnt) > 10000 and h > w and w < width * 0.3 and w > width * 0.15: 222 | expand = 15 223 | tooltip = input_image[y - expand:y + h + expand, x - expand:x + w + expand] 224 | found_tooltip = True 225 | 226 | cv2.rectangle(tooltipcontours_image, (x - expand, y - expand), (x + w + expand, y + h + expand), (0, 255, 0), 3) 227 | else: 228 | cv2.rectangle(tooltipcontours_image, (x, y), (x + w, y + h), (255, 0, 0), 3) 229 | 230 | tooltipcontours_image = cv2.resize(tooltipcontours_image, (0,0), fx = 0.5, fy = 0.5) 231 | else: 232 | tooltip = input_image 233 | 234 | if find_tooltip == True and found_tooltip != True: 235 | print("ERROR: Failed to find tooltip...") 236 | return None, 0.0 237 | else: 238 | # use paddle ocr with custom d4 tooltip trained recognition model 239 | #tooltip = cv2.resize(tooltip, (0,0), fx = 0.5, fy = 0.5) 240 | start_ocr = time.time() 241 | result = self.ocr.ocr(tooltip, cls=False) 242 | end_ocr = time.time() 243 | 244 | if debug == True: 245 | # print ocr results 246 | print('OCR results:') 247 | for idx in range(len(result)): 248 | res = result[idx] 249 | for line in res: 250 | print(line) 251 | print('') # newline 252 | 253 | result = result[0] 254 | boxes = [line[0] for line in result] 255 | txts = [line[1][0] for line in result] 256 | scores = [line[1][1] for line in result] 257 | 258 | # use scale invariant template matching to find all symbols denoting different types of tooltip lines 259 | data, templmatch_image = scaleInvariantMultiTemplateMatch(tooltip, self.templates, show_log=debug) 260 | 261 | if debug == True: 262 | ocr_img = draw_ocr(tooltip, boxes, txts, scores, font_path='C:/Windows/fonts/Arial.ttf') 263 | cv2.imshow('ocr_img', ocr_img) 264 | if find_tooltip == True: 265 | cv2.imshow('tooltip_countours_img', tooltipcontours_image) 266 | cv2.imshow('template_match_img', templmatch_image) 267 | 268 | # ===================================================== 269 | # build item tooltip data object for json serialization 270 | # ===================================================== 271 | lc = LineCollection(result) 272 | 273 | # find item power and split on it (item power always exists in the tooltip) 274 | txt_item_power, before, after = lc.splitat(lc.find('Item Power', strip='^([\s\d\+]+)')) 275 | item_power, item_power_upgraded = txt_item_power.re_get('(?P\d+)(?:\+(?P\d+))?', 'ip', 'ipu', flags=re.I) 276 | 277 | # spit on requires level because nothing below it is of interest 278 | txt_requires_level, after, tmp1 = after.splitat(after.find('Requires Level', strip='([\s\d]+)$')) 279 | 280 | # remove anything after last data entry (socket, aspect or affix) 281 | # strips flavor text from our collection 282 | if (len(data) >= 1): 283 | key, pt1, pt2, max_val = data[-1] 284 | after = after.takewhile(lambda x : x.tl[1] <= pt2[1] or x.tl[0] >= pt2[0] - int((pt2[0] - pt1[0]) / 2), remove=False) 285 | 286 | # item name is always in uppercase 287 | item_name = before.takewhile(lambda x : x.txt.upper() == x.txt) 288 | # item type follows item name and is the last line before item power 289 | item_type = before #before[len(item_name):] 290 | 291 | # upgrades follows item power (if it exists) 292 | txt_upgrades, *_ = lc.splitat(lc.find('Upgrades:', strip='([\s\d/]+)$')) 293 | item_upgrades_current, item_upgrades_max = txt_upgrades.re_get('(?P\d+)/(?P\d+)', 'upc', 'upm') 294 | if txt_upgrades.index != -1: 295 | after.lines.remove(txt_upgrades) 296 | 297 | item = {'affixes': [], 'stats': [], 'sockets': [], 'aspect': None } 298 | item['name'] = lines_join(' ', [l.txt.strip() for l in item_name]) 299 | item['type'] = lines_join(' ', [l.txt.strip() for l in item_type]) 300 | item['item_power'] = item_power 301 | item['item_power_upgraded'] = item_power_upgraded 302 | item['item_upgrades_current'] = item_upgrades_current 303 | item['item_upgrades_max'] = item_upgrades_max 304 | 305 | if debug == True: 306 | print('Line cy offsets:') 307 | for a in after.lines: 308 | print(f'[cy: {a.cy}]: {a.txt}') 309 | print('') # newline 310 | 311 | # fetch lines belonging to any data entries using relative positions 312 | for i, d in enumerate(data): 313 | key_1, pt1_1, pt2_1, max_val_1 = d 314 | 315 | # anything above the first data entry is going to be weapon or armor stats 316 | stats = after.takebetween(0, pt1_1[1]) 317 | for s in stats.lines: 318 | item['stats'].append(lines_join('', [s.txt])) 319 | 320 | # use the position of the next data entry to find the lines belonging to the current data entry 321 | if len(data) > i + 1: 322 | key_2, pt1_2, pt2_2, max_val_2 = data[i + 1] 323 | else: 324 | key_2, pt1_2, pt2_2, max_val_2 = None, [0, 9999], [0, 9999], 0.0 325 | 326 | dline = after.takebetween(pt1_1[1], pt2_2[1] - (pt2_2[1] - pt1_2[1])) 327 | if key_1 == 'affix' or key_1 == 'reroll': 328 | item['affixes'].append(lines_join(' ', [l.txt.strip() for l in dline or after])) 329 | elif key_1 == 'wstat': 330 | item['stats'].append(lines_join(' ', [l.txt.strip() for l in dline or after])) 331 | elif key_1 == 'aspect': 332 | item['aspect'] = lines_join(' ', [l.txt.strip() for l in dline or after]) 333 | elif key_1 == 'socket': 334 | item['sockets'].append(lines_join(' ', [l.txt.strip() for l in dline or after])) 335 | 336 | # serialize to json 337 | jsonstr = json.dumps(item, sort_keys=True, indent=3) 338 | return jsonstr, (end_ocr - start_ocr) 339 | 340 | # use substitution to fix common errors in recognized text 341 | def lines_join(separator, iterable): 342 | result = separator.join(iterable) 343 | 344 | result = re.sub('(?<=[a-z])(?:\s+-\s+|\s+-|-\s+)(?=[a-z])', '-', result, flags=re.I) # [a -b]: missing space 345 | result = re.sub('^0\+', '+', result) # 0+: symbol recognized as 0 346 | result = re.sub('%([a-z])', r'% \1', result, flags=re.I) # %a: missing space 347 | result = re.sub(r'\b([A-Z])(([a-z]+)([A-Z]+)([a-z]*))\b', lambda m: m.group(1) + m.group(2).lower(), result) # PulveriZe: in word case mismatch 348 | result = re.sub('\]([a-z])', r'] \1', result, flags=re.I) # ]A: missing space 349 | result = re.sub('\[([\d\.]+)[\s-]+([\d\.]+)\]%', r'[\1 - \2]%', result) # [15.0 -20.0]%: missing space 350 | result = re.sub('([\d,]+)(?:\s+-\s+|\s+-|-\s+)([\d,]+)', r'\1 - \2', result) # [800 -1,000]: missing space 351 | 352 | return result 353 | 354 | # scale invariant multi template matching using image pyramid 355 | def scaleInvariantMultiTemplateMatch(img_source, img_tmpl_dict, matchTemplateThreshold=0.95, show_log=False): 356 | source = cv2.cvtColor(img_source, cv2.COLOR_BGR2GRAY) 357 | 358 | # pre-process templates 359 | tmpls = [] 360 | tmpls_mask = [] 361 | for k, tmpl in img_tmpl_dict.items(): 362 | if isinstance(tmpl, list): 363 | tmpls.append(cv2.cvtColor(tmpl[0], cv2.COLOR_BGR2GRAY)) 364 | tmpls_mask.append(cv2.cvtColor(tmpl[1], cv2.COLOR_BGR2GRAY)) 365 | else: 366 | tmpls.append(cv2.cvtColor(tmpl, cv2.COLOR_BGR2GRAY)) 367 | tmpls_mask.append(None) 368 | 369 | # per-downsampled source results 370 | results = [] 371 | best_index = -1 372 | 373 | # template matching using 20 downsampled sources in scale range (0.2 - 1.0) 374 | for i, scale in enumerate(np.linspace(0.2, 1.0, 20)[::-1]): 375 | resized = cv2.resize(source, (0, 0), fx = scale, fy = scale) 376 | ratio = source.shape[1] / float(resized.shape[1]) 377 | 378 | # per-template results 379 | t_results = {} 380 | max_value_high = 0 381 | 382 | # template match each input template 383 | for j, (k, img_tmpl) in enumerate(img_tmpl_dict.items()): 384 | tmpl = tmpls[j] 385 | tmpl_mask = tmpls_mask[j] 386 | 387 | # skip to next iteration if downsampled source is smaller than current template 388 | if resized.shape[0] < tmpl.shape[0] or resized.shape[1] < tmpl.shape[1]: 389 | t_results[k] = (None, None, 0.0, (0, 0), ratio) 390 | continue 391 | 392 | # match template and threshold with custom threshold 393 | result = cv2.matchTemplate(resized, tmpl, cv2.TM_CCORR_NORMED, mask=tmpl_mask) 394 | T, threshed = cv2.threshold(result, matchTemplateThreshold, 1., cv2.THRESH_TOZERO) 395 | 396 | # find best match score for current downsampled source in effort 397 | # to figure out which scale give the best overall results 398 | min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(threshed) 399 | 400 | t_results[k] = (result, threshed, max_val, max_loc) 401 | if max_val > max_value_high: 402 | max_value_high = max_val 403 | 404 | results.append((t_results, max_value_high, ratio)) 405 | 406 | if best_index == -1: 407 | best_index = 0 408 | elif max_value_high > results[best_index][1]: 409 | best_index = i 410 | 411 | # best downsampled scale 412 | t_results, max_value_high, ratio = results[best_index] 413 | if show_log == True: 414 | print('Template matching:') 415 | print(f'best_index: {best_index}, r: {ratio}', end='\n\n') 416 | 417 | dest = img_source.copy() 418 | data = [] 419 | 420 | # find best template match locations and draw them on a copy of source 421 | if show_log == True: 422 | print('Found templates:') 423 | 424 | for i, (k, img_tmpl) in enumerate(img_tmpl_dict.items()): 425 | h, w = tmpls[i].shape 426 | result, threshed, max_val, max_loc = t_results[k] 427 | 428 | while max_val > 0.9: 429 | pt1 = (int(round(max_loc[0] * ratio, 0)), int(round(max_loc[1] * ratio, 0))) 430 | pt2 = (int(round((max_loc[0] + w + 1) * ratio, 0)), int(round((max_loc[1] + h + 1) * ratio, 0))) 431 | 432 | if show_log == True: 433 | print(f'key: {k}, region: {pt1} => {pt2}, max_val: {max_val}') 434 | 435 | data.append((k, pt1, pt2, max_val)) 436 | cv2.rectangle(dest, pt1, pt2, (0,255,0), 2) 437 | 438 | threshed[max_loc[1] - h // 2 : max_loc[1] + h // 2 + 1, max_loc[0] - w // 2 : max_loc[0] + w // 2 + 1] = 0 439 | min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(threshed) 440 | 441 | if show_log == True and len(data) == 0: 442 | print('[None]', end='\n\n') 443 | elif show_log == True: 444 | print('') # newline 445 | 446 | # sort locations on y-offset in ascending order 447 | data.sort(key=lambda x : x[1][1]) 448 | 449 | return data, dest 450 | 451 | # create a tkinter window for viewing item tooltip data and 452 | # navigating between images in the current directory 453 | def showItemDataFrame(source_path, find_tooltip=True): 454 | dnav = SimpleDirectoryNavigator(source_path) 455 | it_ocr = D4ItemTooltipOCR() 456 | 457 | def toggle_button(btn, state): 458 | if state == True: 459 | btn["state"] = tk.NORMAL 460 | else: 461 | btn["state"] = tk.DISABLED 462 | 463 | def showImage(offsetFromCurrent = 0): 464 | image_path, has_prev, has_next = dnav.getImagePath(offsetFromCurrent) 465 | 466 | toggle_button(prev_button, has_prev) 467 | toggle_button(next_button, has_next) 468 | 469 | if image_path: 470 | jsonstr = it_ocr.processImage(image_path, find_tooltip, debug=True) 471 | 472 | win.title(f'D4 Item Data: {dnav.current_image_filename}') 473 | T.delete(1.0, tk.END) 474 | 475 | if jsonstr: 476 | T.insert(tk.END, jsonstr) 477 | else: 478 | win.title(f'D4 Item Data') 479 | T.delete(1.0, tk.END) 480 | 481 | def saveJson(): 482 | filename_without_ext = os.path.splitext(dnav.current_image_filename)[0] 483 | json_path = os.path.join(dnav.dir_path, filename_without_ext + '.json') 484 | print(f'Saving: \'{json_path}\'') 485 | 486 | jsonstr = T.get(1.0, tk.END) 487 | with open(json_path, 'w') as outfile: 488 | print('Writing:') 489 | print(jsonstr) 490 | outfile.write(jsonstr) 491 | 492 | win = tk.Tk() 493 | win.geometry("800x600") 494 | win.title("D4 Item Data") 495 | 496 | button_frame = tk.Frame(win) 497 | button_frame.pack(fill=tk.X, side=tk.BOTTOM) 498 | 499 | button_frame_top = tk.Frame(win) 500 | button_frame_top.pack(fill=tk.X, side=tk.TOP) 501 | 502 | scroll_v = tk.Scrollbar(win) 503 | scroll_v.pack(side=tk.RIGHT,fill="y") 504 | scroll_h = tk.Scrollbar(win, orient=tk.HORIZONTAL) 505 | scroll_h.pack(side=tk.BOTTOM, fill= "x") 506 | 507 | T = tk.Text(win, height = 50, width = 500, yscrollcommand= scroll_v.set,xscrollcommand = scroll_h.set, wrap=tk.NONE,) 508 | T.config(font=("Courier New", 11)) 509 | T.pack(fill=tk.BOTH, expand=0) 510 | 511 | scroll_h.config(command = T.xview) 512 | scroll_v.config(command = T.yview) 513 | 514 | prev_button = tk.Button(button_frame, text="Prev", command = lambda:showImage(-1)) 515 | next_button = tk.Button(button_frame, text="Next", command = lambda:showImage(1)) 516 | 517 | button_frame.columnconfigure(0, weight=1) 518 | button_frame.columnconfigure(1, weight=1) 519 | 520 | prev_button.grid(row=0, column=0, sticky=tk.W+tk.E) 521 | next_button.grid(row=0, column=1, sticky=tk.W+tk.E) 522 | 523 | status_label = tk.Label(button_frame_top, text="Status") 524 | save_button = tk.Button(button_frame_top, text="Save JSON", command = saveJson) 525 | 526 | button_frame_top.columnconfigure(0, weight=4) 527 | button_frame_top.columnconfigure(1, weight=1) 528 | 529 | status_label.grid(row=0, column=0, sticky=tk.W+tk.E) 530 | save_button.grid(row=0, column=1, sticky=tk.W+tk.E) 531 | 532 | win.after(0, showImage) 533 | tk.mainloop() 534 | 535 | # run tests in directory (checks all images that have a definition (.json)) 536 | def runTestsInDir(dir_path): 537 | dnav = SimpleDirectoryNavigator(dir_path, run_tests_mode=True) 538 | it_ocr = D4ItemTooltipOCR() 539 | 540 | print("THIS FEATURE IS NOT FULLY IMPLEMENTED") 541 | 542 | while (t := dnav.getImagePath_RunTestsMode(1)) is not None: 543 | path, has_prev, has_next = t 544 | print(f"Test path: {path}") 545 | 546 | print("") 547 | 548 | 549 | if __name__ == '__main__': 550 | parser = argparse.ArgumentParser() 551 | parser.add_argument('--source-img', type=str, default='examples/screenshot_001.png', help='path to the source image') 552 | parser.add_argument('--json-output', type=str, default=None, help='output path for item tooltip json data') 553 | parser.add_argument('--find-tooltip', default=True, type=lambda x: x.lower() not in ['false', 'no', '0', 'None'], help='toggle find tooltip in source [true/false]') 554 | parser.add_argument('--debug', default=False, type=lambda x: x.lower() not in ['false', 'no', '0', 'None'], help='toggle debug mode [true/false]') 555 | parser.add_argument('--run-tests-dir', type=str, default=None, help='run tests in directory') 556 | opt = parser.parse_args() 557 | 558 | print("") 559 | print("======================================") 560 | print("Diablo IV: Item Tooltip OCR") 561 | print("======================================") 562 | print("") 563 | 564 | # check that source image exists 565 | if os.path.isfile(opt.source_img) == False: 566 | print('ERROR: Source image not found...', end='\n\n') 567 | parser.print_help() 568 | exit(1) 569 | 570 | # use debug mode 571 | if opt.debug == True: 572 | showItemDataFrame(opt.source_img, opt.find_tooltip) 573 | cv2.destroyAllWindows() 574 | 575 | # run tests mode 576 | elif opt.run_tests_dir != None: 577 | runTestsInDir(opt.run_tests_dir) 578 | cv2.destroyAllWindows() 579 | 580 | # output json data from source 581 | else: 582 | it_ocr = D4ItemTooltipOCR() 583 | jsonstr = it_ocr.processImage(opt.source_img, opt.find_tooltip, debug=False) 584 | 585 | if (opt.json_output): 586 | with open(opt.json_output, 'w') as outfile: 587 | outfile.write(jsonstr) 588 | print(f'Item tooltip json written to \'{opt.json_output}\'.', end='\n\n') 589 | else: 590 | print(jsonstr) -------------------------------------------------------------------------------- /examples/screenshot_001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/examples/screenshot_001.png -------------------------------------------------------------------------------- /examples/tooltip_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/examples/tooltip_001.jpg -------------------------------------------------------------------------------- /paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdiparams: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdiparams -------------------------------------------------------------------------------- /paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdiparams.info: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdiparams.info -------------------------------------------------------------------------------- /paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdmodel: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/paddleocr-models/en_PP-OCRv3_rec-d4_tooltip/inference.pdmodel -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.23.5 2 | opencv_contrib_python==4.6.0.66 3 | opencv_python==4.6.0.66 4 | opencv_python_headless==4.6.0.66 5 | paddlepaddle==2.5.0 6 | paddleocr==2.6.1.3 7 | python_Levenshtein==0.21.1 8 | -------------------------------------------------------------------------------- /templates/affix.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/affix.png -------------------------------------------------------------------------------- /templates/enchanted_rerolled.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/enchanted_rerolled.png -------------------------------------------------------------------------------- /templates/inprint_aspect.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/inprint_aspect.png -------------------------------------------------------------------------------- /templates/socket.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/socket.png -------------------------------------------------------------------------------- /templates/socket_mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/socket_mask.png -------------------------------------------------------------------------------- /templates/socket_mask_new.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/socket_mask_new.png -------------------------------------------------------------------------------- /templates/weapon_stat.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mxtsdev/d4-item-tooltip-ocr/493ac5697fd4a2a5c0074724fd76e0873a914b59/templates/weapon_stat.png --------------------------------------------------------------------------------