From c8eb88f67201a02d3a58eb0067f2aaffe6c84ce5 Mon Sep 17 00:00:00 2001 From: mgleed Date: Mon, 21 Feb 2022 19:54:29 -0500 Subject: [PATCH] Feature: OCR for item drops and item description boxes (#570) * initial * bugfixes. now stable & testing * narrow ROI to search for 'CLICK TO' tooltip --- assets/templates/ui/inventory/to_tooltip.png | Bin 0 -> 2159 bytes assets/tessdata/ocr_errors.csv | 1 + config/game.ini | 1 + config/params.ini | 1 + src/config.py | 3 +- src/game_stats.py | 4 +- src/item/item_cropper.py | 77 +++++++++++++++++-- src/item/item_finder.py | 9 ++- src/item/pickit.py | 23 ++++++ src/messages/discord_embeds.py | 5 +- src/messages/generic_api.py | 10 +-- src/messages/messenger.py | 4 +- src/ocr.py | 26 ++++++- src/template_finder.py | 1 - src/ui_components/inventory.py | 23 ++++++ src/utils/gen_ocr_samples.py | 74 ++++++++++++++++++ 16 files changed, 237 insertions(+), 25 deletions(-) create mode 100644 assets/templates/ui/inventory/to_tooltip.png create mode 100644 src/utils/gen_ocr_samples.py diff --git a/assets/templates/ui/inventory/to_tooltip.png b/assets/templates/ui/inventory/to_tooltip.png new file mode 100644 index 0000000000000000000000000000000000000000..2dcfcb4db8f132badf0799beb9ab0869f8a8cc54 GIT binary patch literal 2159 zcmcImeNYr-99~h=RM0{lCR4YCNl^E8Z|~slcDM{Ulq(Mr&IHVGZ}+{s1#b7eT{ytc zQE>u8r8E^8h*5KJj4~uYC^FFuCon%kQUh#A8d0&*qBQJ1I4-AA%Oq6H{&98QQvkjPm%PKM0@GHL-fUeIIRwVz`s&*`y+;U?1Li~)8& zF3$y$^5V^GUIwea~zq1MCa(SK3*EmJ+BE2?Sse}dMwIg5KT5E zpfREgpc;h|XGxMoH9Ccbv}!FtHxs2uN=Yawg4E(l9j#Q-6oq;(44QLsRyxrz+iMK2 z^q5_ioistX-EM_jr4U^$`xcL2N$H^0R{t*SQqb7x}SZ=7Jnh#baAC>5l~rvaeJ^kk^R;X_MuY&Nvw)%!?d$(-v1L?V#H zEEfx6vwNs}6yQsHDcO_Kdk z6D$Odk(t4PMUJ=T4hCv6(MCa%8G!{xgC2vT6g(h9q&8Ra3Z*(J0}K zi%@YoEgQj5eY)UzgUDui!sj{PU$>mdLJnVqXf4obEIJi~Q^2CcIi;G#84JbXTm;K# zS(VkIW0(Q9c`hD)I!xxntR7VyWK^odwQ7~sitE^LHLi?Mb2vk4DV)(!Dh-HWRXPr; z;#k@$x*QCw6W+nt0O1sD82XQz7%@|HnM4l8T;-`|$PT@T<0Y87T<`ah1eSVdnLOH8 zbTq?yYF&@9o+JYf^S8+x|P*!Ti<2qlWytVeRXCrqwq7 zP!P5GWJxT7-X#AZe<}KS^??Mz_Nx|dphZ+d_4(6vp4^MbnGyydu@xM=cn&7JUG zi&JrAq~DU2%SN>_t|&H6 ze>ru}45T6EnAoEU!mus&kUJg1yxbjWhPd@zGbSGzZ3`&7*-Bs)M+>W(vY#oK9lc2? zSNLz)9FiYmET1t4IaYpfN7HuU)R4Vhs$=w4v-AGiv@flJ=V$FI9(T3% ztvotu%({uMTXRXIy?1L;)9JG@j^^D5Dz4Rc4!dwV{>a;lC%mzEd`V?KKsxiLPV3xu z{O*Jizt!A1)#i6N`{$Mu9i=zRYm-037GGa5C7ovL9;jE|zpX&N+s~K0xcUQA&CcTs zf_I&tVyibczf*&^T9mgUmuLNYuHn-4_S-+1CHDK`mW@?S`itHEGZnVT*n8^^w!Yxl zGu`NSV56OlY0R^4-!)HJdY5BzHyC9WzwyD}26kUuv5D?AVOE>o#-*v=cX~&k4K!@(cC+ O5*TCS4OP)g*8C0iRT9Dg literal 0 HcmV?d00001 diff --git a/assets/tessdata/ocr_errors.csv b/assets/tessdata/ocr_errors.csv index d322639..3c7dff3 100644 --- a/assets/tessdata/ocr_errors.csv +++ b/assets/tessdata/ocr_errors.csv @@ -1,6 +1,7 @@ SHIFLD, SHIELD SPFAR, SPEAR GLOVFS, GLOVES +GOLP, GOLD TELEFORT, TELEPORT TROPHV, TROPHY CLAVMORE, CLAYMORE diff --git a/config/game.ini b/config/game.ini index c1d102a..abee9df 100644 --- a/config/game.ini +++ b/config/game.ini @@ -2,6 +2,7 @@ ; min and max hsv range (opencv format: h: [0-180], s: [0-255], v: [0, 255]) ; h_min, s_min, v_min, h_max, s_max, v_max black=0,0,0,180,255,15 +black_descr=0,0,0,180,255,25 item_highlight=90,235,130,115,255,160 white=0,0,150,180,20,255 gray=0,0,90,180,20,130 diff --git a/config/params.ini b/config/params.ini index a0c2f3e..3acb08b 100644 --- a/config/params.ini +++ b/config/params.ini @@ -255,3 +255,4 @@ hwnd_window_title= hwnd_window_process=D2R\.exe ;If you want to control Hyper-V window from host use 0,51 here window_client_area_offset=0,0 +use_ocr=0 diff --git a/src/config.py b/src/config.py index 1695c56..2f99f14 100644 --- a/src/config.py +++ b/src/config.py @@ -316,7 +316,8 @@ class Config: "graphic_debugger_key": self._select_val("advanced_options", "graphic_debugger_key"), "hwnd_window_title": _default_iff(Config()._select_val("advanced_options", "hwnd_window_title"), ''), "hwnd_window_process": _default_iff(Config()._select_val("advanced_options", "hwnd_window_process"), ''), - "window_client_area_offset": tuple(map(int, Config()._select_val("advanced_options", "window_client_area_offset").split(","))) + "window_client_area_offset": tuple(map(int, Config()._select_val("advanced_options", "window_client_area_offset").split(","))), + "use_ocr": bool(int(self._select_val("advanced_options", "use_ocr"))), } self.items = {} diff --git a/src/game_stats.py b/src/game_stats.py index 8d5eb03..4ee0c3a 100644 --- a/src/game_stats.py +++ b/src/game_stats.py @@ -41,7 +41,7 @@ class GameStats: if self._location not in self._location_stats: self._location_stats[self._location] = { "items": [], "deaths": 0, "chickens": 0, "merc_deaths": 0, "failed_runs": 0 } - def log_item_keep(self, item_name: str, send_message: bool, img: np.ndarray): + def log_item_keep(self, item_name: str, send_message: bool, img: np.ndarray, ocr_text: str = None): Logger.debug(f"Stashed and logged: {item_name}") filtered_items = ["_potion", "misc_gold"] if self._location is not None and not any(substring in item_name for substring in filtered_items): @@ -49,7 +49,7 @@ class GameStats: self._location_stats["totals"]["items"] += 1 if send_message: - self._messenger.send_item(item_name, img, self._location) + self._messenger.send_item(item_name, img, self._location, ocr_text) def log_death(self, img: str): self._death_counter += 1 diff --git a/src/item/item_cropper.py b/src/item/item_cropper.py index 82b8e55..8d16419 100644 --- a/src/item/item_cropper.py +++ b/src/item/item_cropper.py @@ -4,20 +4,31 @@ from dataclasses import dataclass import time from utils.misc import color_filter, erode_to_black +from template_finder import TemplateFinder +from ocr import Ocr, OcrResult from config import Config +from logger import Logger # TODO: With OCR we can then add a "text" field to this class @dataclass class ItemText: - color_key: str = None + color: str = None roi: list[int] = None data: np.ndarray = None + ocr_result: OcrResult = None + clean_img: np.ndarray = None + def __getitem__(self, key): + return super().__getattribute__(key) class ItemCropper: def __init__(self): + self._ocr = Ocr() + self._gaus_filter = (19, 1) self._expected_height_range = [round(num) for num in [x / 1.5 for x in [14, 40]]] self._expected_width_range = [round(num) for num in [x / 1.5 for x in [60, 1280]]] + self._box_expected_width_range=[200, 900] + self._box_expected_height_range=[24, 710] self._hud_mask = cv2.imread(f"assets/hud_mask.png", cv2.IMREAD_GRAYSCALE) self._hud_mask = cv2.threshold(self._hud_mask, 1, 255, cv2.THRESH_BINARY)[1] @@ -70,28 +81,80 @@ class ItemCropper: max_idx = color_averages.index(max(color_averages)) if key == self._item_colors[max_idx]: item_clusters.append(ItemText( - color_key = key, + color = key, roi = [x, y, w, h], - data = cropped_item + data = cropped_item, + clean_img = cleaned_img[y:y+h, x:x+w] )) debug_str += f" | cluster: {time.time() - start}" # print(debug_str) + if Config().advanced_options["use_ocr"]: + cluster_images = [ key["clean_img"] for key in item_clusters ] + results = self._ocr.image_to_text(cluster_images, model = "engd2r_inv_th_fast", psm = 7) + for count, cluster in enumerate(item_clusters): + setattr(cluster, "ocr_result", results[count]) return item_clusters + def crop_item_descr(self, inp_img: np.ndarray, all_results: bool = False, inventory_side: str = "right") -> ItemText: + """ + Crops visible item description boxes / tooltips + :inp_img: image from hover over item of interest. + :param all_results: whether to return all possible results (True) or the first result (False) + :inventory_side: enter either "left" for stash/vendor region or "right" for user inventory region + """ + results=[] + black_mask, _ = color_filter(inp_img, Config().colors["black"]) + contours = cv2.findContours(black_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + contours = contours[0] if len(contours) == 2 else contours[1] + for cntr in contours: + x, y, w, h = cv2.boundingRect(cntr) + cropped_item = inp_img[y:y+h, x:x+w] + avg = np.average(cv2.cvtColor(cropped_item, cv2.COLOR_BGR2GRAY)) + mostly_dark = True if 0 < avg < 20 else False + contains_black = True if np.min(cropped_item) < 14 else False + contains_white = True if np.max(cropped_item) > 250 else False + contains_orange = False + if not contains_white: + #check for orange (like key of destruction, etc.) + orange_mask, _ = color_filter(cropped_item, Config().colors["orange"]) + contains_orange = np.min(orange_mask) > 0 + expected_height = True if (self._box_expected_height_range[0] < h < self._box_expected_height_range[1]) else False + expected_width = True if (self._box_expected_width_range[0] < w < self._box_expected_width_range[1]) else False + box2 = Config().ui_roi[f"{inventory_side}_inventory"] + overlaps_inventory = False if (x+w 720 else 35 + found_footer = TemplateFinder().search(["TO_TOOLTIP"], inp_img, threshold=0.8, roi=[x, y+h, w, footer_height_max]).valid + if found_footer: + ocr_result = None + if Config().advanced_options["use_ocr"]: + ocr_result = self._ocr.image_to_text(cropped_item, psm=6)[0] + results.append(ItemText( + color = "black", + roi = [x, y, w, h], + data = cropped_item, + ocr_result = ocr_result + )) + if not all_results: + break + return results if __name__ == "__main__": import keyboard import os from screen import grab + from template_finder import TemplateFinder keyboard.add_hotkey('f12', lambda: os._exit(1)) cropper = ItemCropper() while 1: img = grab().copy() - res = cropper.crop(img) - for cluster in res: - x, y, w, h = cluster.roi - cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 1) + results = cropper.crop_item_descr(img, all_results=True, ocr=False) + for res in results: + if res["color"]: + x, y, w, h = res.roi + cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 1) + Logger.debug(f"{res.ocr_result['text']}") cv2.imshow("res", img) cv2.waitKey(1) diff --git a/src/item/item_finder.py b/src/item/item_finder.py index f874e08..e8c5928 100644 --- a/src/item/item_finder.py +++ b/src/item/item_finder.py @@ -9,6 +9,8 @@ import math from config import Config from utils.misc import color_filter, cut_roi from item import ItemCropper +from template_finder import TemplateFinder +from ocr import OcrResult @dataclass @@ -24,6 +26,8 @@ class Item: score: float = -1.0 dist: float = -1.0 roi: list[int] = None + color: str = None + ocr_result: OcrResult = None def __getitem__(self, key): return super().__getattribute__(key) @@ -119,6 +123,8 @@ class ItemFinder: item.roi = [max_loc[0] + x, max_loc[1] + y, template.data.shape[1], template.data.shape[0]] center_abs = (item.center[0] - (inp_img.shape[1] // 2), item.center[1] - (inp_img.shape[0] // 2)) item.dist = math.dist(center_abs, (0, 0)) + item.ocr_result = cluster.ocr_result + item.color = cluster.color if item is not None and self._items_to_pick[item.name].pickit_type: item_list.append(item) elapsed = time.time() - start @@ -130,6 +136,7 @@ class ItemFinder: if __name__ == "__main__": from screen import grab from config import Config + item_finder = ItemFinder() while 1: # img = cv2.imread("") @@ -139,7 +146,7 @@ if __name__ == "__main__": # print(item.name + " " + str(item.score)) cv2.circle(img, item.center, 5, (255, 0, 255), thickness=3) cv2.rectangle(img, item.roi[:2], (item.roi[0] + item.roi[2], item.roi[1] + item.roi[3]), (0, 0, 255), 1) - # cv2.putText(img, item.name, item.center, cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1, cv2.LINE_AA) + cv2.putText(img, item.ocr_result["text"], item.center, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA) # img = cv2.resize(img, None, fx=0.5, fy=0.5) cv2.imshow('test', img) cv2.waitKey(1) \ No newline at end of file diff --git a/src/item/pickit.py b/src/item/pickit.py index 25c26df..90a9f9d 100644 --- a/src/item/pickit.py +++ b/src/item/pickit.py @@ -40,6 +40,8 @@ class PickIt: curr_item_to_pick: Item = None same_item_timer = None did_force_move = False + done_ocr=False + while not time_out: if (time.time() - start) > 28: time_out = True @@ -48,6 +50,23 @@ class PickIt: img = grab() item_list = self._item_finder.search(img) + if Config().advanced_options["use_ocr"] and not done_ocr: + timestamp = time.strftime("%Y%m%d_%H%M%S") + for cnt, item in enumerate(item_list): + for cnt2, x in enumerate(item.ocr_result['word_confidences']): + found_low_confidence = False + if x <= 88: + try: + Logger.debug(f"Low confidence word #{cnt2}: {item.ocr_result['original_text'].split()[cnt2]} -> {item.ocr_result['text'].split()[cnt2]}, Conf: {x}, save screenshot") + found_low_confidence = True + except: pass + if found_low_confidence and Config().general["loot_screenshots"]: + cv2.imwrite(f"./loot_screenshots/ocr_drop_{timestamp}_{cnt}_o.png", item.ocr_result['original_img']) + cv2.imwrite(f"./loot_screenshots/ocr_drop_{timestamp}_{cnt}_n.png", item.ocr_result['processed_img']) + with open(f"./loot_screenshots/ocr_drop_{timestamp}_{cnt}_o.gt.txt", 'w') as f: + f.write(item.ocr_result['text']) + done_ocr = True + # Check if we need to pick up certain pots more pots need_pots = belt.get_pot_needs() if need_pots["mana"] <= 0: @@ -110,6 +129,10 @@ class PickIt: # no need to stash potions, scrolls, or gold if "potion" not in closest_item.name and "tp_scroll" != closest_item.name and "misc_gold" not in closest_item.name: found_items = True + if Config().advanced_options["use_ocr"]: + for item in item_list: + Logger.debug(f"OCR DROP: Name: {item.ocr_result['text']}, Conf: {item.ocr_result['word_confidences']}") + prev_cast_start = char.pick_up_item((x_m, y_m), item_name=closest_item.name, prev_cast_start=prev_cast_start) if not char.capabilities.can_teleport_natively: diff --git a/src/messages/discord_embeds.py b/src/messages/discord_embeds.py index 3692a72..ac3b3a1 100644 --- a/src/messages/discord_embeds.py +++ b/src/messages/discord_embeds.py @@ -18,7 +18,7 @@ class DiscordEmbeds(GenericApi): except InvalidArgument: Logger.warning(f"Your custom_message_hook URL {Config().general['custom_message_hook']} is invalid, Discord updates will not be sent") - def send_item(self, item: str, image: np.ndarray, location: str): + def send_item(self, item: str, image: np.ndarray, location: str, ocr_text: str = None): imgName = item.replace('_', '-') _, w, _ = image.shape @@ -32,6 +32,7 @@ class DiscordEmbeds(GenericApi): ) e.set_thumbnail(url=f"{self._psnURL}41L6bd712.png") e.set_image(url=f"attachment://{imgName}.png") + e.add_field(name="OCR Text", value=f"{ocr_text}", inline=False) self._send_embed(e, file) def send_death(self, location, image_path): @@ -96,7 +97,7 @@ class DiscordEmbeds(GenericApi): return Color.blue() def _add_file(self, image_path, image_name): - try: + try: return discord.File(image_path, filename=image_name) except: traceback.print_exc() diff --git a/src/messages/generic_api.py b/src/messages/generic_api.py index 1f3e805..622026b 100644 --- a/src/messages/generic_api.py +++ b/src/messages/generic_api.py @@ -6,18 +6,18 @@ import requests class GenericApi: - def send_item(self, item: str, image: np.ndarray, location: str): + def send_item(self, item: str, image: np.ndarray, location: str, ocr_text: str = None): msg = f"Found {item} at {location}" self._send(msg) - + def send_death(self, location: str, image_path: str = None): msg = f"You have died at {location}" self._send(msg) - + def send_chicken(self, location: str, image_path: str = None): msg = f"You have chickened at {location}" self._send(msg) - + def send_gold(self): msg = f"All stash tabs and character are full of gold, turn of gold pickup" self._send(msg) @@ -31,7 +31,7 @@ class GenericApi: def _send(self, msg: str): msg = f"{Config().general['name']}: {msg}" - + url = Config().general['custom_message_hook'] if not url: return diff --git a/src/messages/messenger.py b/src/messages/messenger.py index c2f3f3d..ec09638 100644 --- a/src/messages/messenger.py +++ b/src/messages/messenger.py @@ -15,8 +15,8 @@ class Messenger: else: self._message_api = None - def send_item(self, item: str, image: np.ndarray, location: str): - self._message_api.send_item(item, image, location) + def send_item(self, item: str, image: np.ndarray, location: str, ocr_text: str = None): + self._message_api.send_item(item, image, location, ocr_text) def send_death(self, location: str, image_path: str = None): self._message_api.send_death(location, image_path) diff --git a/src/ocr.py b/src/ocr.py index 0bd7a92..b46544c 100644 --- a/src/ocr.py +++ b/src/ocr.py @@ -1,3 +1,4 @@ +from concurrent.futures import process from fileinput import close from tesserocr import PyTessBaseAPI, PSM, OEM import numpy as np @@ -119,10 +120,9 @@ class Ocr: if crop_pad: processed_img = self._crop_pad(processed_img) image_is_binary = (image.shape[2] if len(image.shape) == 3 else 1) == 1 and image.dtype == bool - if image_is_binary: - if threshold: - processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY) - processed_img = cv2.threshold(processed_img, threshold, 255, cv2.THRESH_BINARY)[1] + if not image_is_binary and threshold: + processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY) + processed_img = cv2.threshold(processed_img, threshold, 255, cv2.THRESH_BINARY)[1] if invert: if threshold or image_is_binary: processed_img = cv2.bitwise_not(processed_img) @@ -239,6 +239,24 @@ class Ocr: continue break + # case: a solitary S; e.g., " 1 TO S DEFENSE" + cnt=0 + while True: + cnt += 1 + if cnt >30: + Logger.error(f"Error ' S ' -> ' 5 ' on {ocr_output}") + break + if " S " in text: + text = text.replace(" S ", " 5 ") + continue + elif ' I\n' in text: + text = text.replace(' S\n', ' 5\n') + continue + elif '\nI ' in text: + text = text.replace('\nS ', '\n5 ') + continue + break + # case: consecutive I's; e.g., "DEFENSE: II" repeat=False cnt=0 diff --git a/src/template_finder.py b/src/template_finder.py index f8544e1..e479737 100644 --- a/src/template_finder.py +++ b/src/template_finder.py @@ -1,7 +1,6 @@ import cv2 import threading from copy import deepcopy -from item.item_finder import Template from screen import convert_screen_to_monitor, grab from typing import Union from dataclasses import dataclass diff --git a/src/ui_components/inventory.py b/src/ui_components/inventory.py index d03dc1c..979f5bc 100644 --- a/src/ui_components/inventory.py +++ b/src/ui_components/inventory.py @@ -16,6 +16,7 @@ from ui_components import stash from ui_components.stash import gold_full from ui.ui_manager import detect_screen_object, messenger, game_stats, wait_for_screen_object, ScreenObjects from messages import Messenger +from item import ItemCropper messanger = Messenger() @@ -180,6 +181,28 @@ def keep_item(item_finder: ItemFinder, img: np.ndarray, do_logging: bool = True) """ wait(0.2, 0.3) _, w, _ = img.shape + + if Config().advanced_options["use_ocr"]: + item_box = ItemCropper().crop_item_descr(inp_img=img) + if item_box: + item_box = item_box[0] + Logger.debug(f"OCR ITEM DESCR: Mean conf: {item_box.ocr_result.mean_confidence}") + for i, line in enumerate(list(filter(None, item_box.ocr_result.text.splitlines()))): + Logger.debug(f"OCR LINE{i}: {line}") + if Config().general["loot_screenshots"]: + timestamp = time.strftime("%Y%m%d_%H%M%S") + found_low_confidence = False + for cnt, x in enumerate(item_box.ocr_result['word_confidences']): + if x <= 88: + try: + Logger.debug(f"Low confidence word #{cnt}: {item_box.ocr_result['original_text'].split()[cnt]} -> {item_box.ocr_result['text'].split()[cnt]}, Conf: {x}, save screenshot") + found_low_confidence = True + except: pass + if found_low_confidence: + cv2.imwrite(f"./loot_screenshots/ocr_box_{timestamp}_o.png", item_box.ocr_result['original_img']) + cv2.imwrite(f"./loot_screenshots/ocr_box_{timestamp}_n.png", item_box.ocr_result['processed_img']) + + img = img[:, (w//2):,:] original_list = item_finder.search(img) filtered_list = [] diff --git a/src/utils/gen_ocr_samples.py b/src/utils/gen_ocr_samples.py new file mode 100644 index 0000000..1ec3314 --- /dev/null +++ b/src/utils/gen_ocr_samples.py @@ -0,0 +1,74 @@ +from screen import Screen +import cv2 +from config import Config +from utils.misc import cut_roi +import mouse +import keyboard +import os +import time +from screen import grab, convert_monitor_to_screen + +class GenOcrTruth: + def __init__(self): + if not os.path.exists("generated"): + os.system("mkdir generated") + os.system(f"cd generated && mkdir ground-truth") + self._half_width = Config().ui_pos["screen_width"] // 2 + self._half_height = Config().ui_pos["screen_height"] // 2 + self._upper_left = None + + def hook(self, e): + if e.event_type == "down": + if e.name == "f12": + os._exit(1) + img = grab() + loc_monitor = mouse.get_position() + loc_screen = convert_monitor_to_screen(loc_monitor) + if e.name == "f8": + # start template + if self._upper_left is None: + self._upper_left = loc_screen + print(f"stored upper left: {self._upper_left}") + print("Select bottom-right corner of template to create and press f8") + return + # finish template + else: + bottom_right = loc_screen + if bottom_right == self._upper_left: + print(f"roi failed, try again") + self._upper_left = None + return + print(f"stored bottom_right: {bottom_right}") + try: + width = (bottom_right[0] - self._upper_left[0]) + height = (bottom_right[1] - self._upper_left[1]) + template_img = cut_roi(img, [*self._upper_left, width, height]) + basename = f"generated/ground-truth/{time.strftime('%Y%m%d_%H%M%S')}" + cv2.imshow(time.strftime('%Y%m%d_%H%M%S'), template_img) + cv2.waitKey(1) + print(f"new template: {basename} = ") + print(f"Enter true text:") + truth = input() + if truth: + with open(f"{basename}.gt.txt", 'w') as f: + f.write(truth) + cv2.imwrite(f"{basename}.png", template_img) + cv2.destroyAllWindows() + cv2.waitKey(1) + print(f"saved {basename}") + else: + print(f"skipped {basename}") + except: + print(f"template save failed, try again") + self._upper_left = None + print("Select top-left corner of template to create and press f8 | F12 to exit") + return + else: + return + +if __name__ == "__main__": + keyboard.add_hotkey('f12', lambda: print('Force Exit (f12)') or os._exit(1)) + recorder = GenOcrTruth() + print("Select top-left corner of template to create and press f8 | F12 to exit") + keyboard.hook(recorder.hook, suppress=True) + while True: pass \ No newline at end of file