Skip to content

Commit

Permalink
chore: 整合OCR相关函数
Browse files Browse the repository at this point in the history
  • Loading branch information
@weiduhuo @Night-stars-1
weiduhuo authored and Night-stars-1 committed Oct 12, 2023
1 parent 2ffb0b3 commit e50983d
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 44 deletions.
42 changes: 10 additions & 32 deletions utils/calculated.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -607,7 +607,7 @@ def ocr_pos(self, characters:str = None, points = (0,0,0,0)):
#img_fp = self.take_screenshot(points)
#out = self.ocr.ocr(img_fp)
#data = {i['text']: i['position'] for i in out}
data = self.part_ocr_other(points)
data = self.part_ocr(points)
log.debug(data)
if not characters:
characters = list(data.keys())[0]
Expand All @@ -617,12 +617,10 @@ def ocr_pos(self, characters:str = None, points = (0,0,0,0)):
pos = data[characters] if characters in data else None
return characters, pos

def ocr_pos_sp(self, characters_list:list[str] = None, points = (0,0,0,0), number = False, debug = False, img_pk:tuple = None) -> Union[int, str]:
def ocr_pos_for_singleLine(self, characters_list:list[str] = None, points = (0,0,0,0), number = False, debug = False, img_pk:tuple = None) -> Union[int, str]:
"""
说明：
或取指定坐标的文字
当characters_list为空，直接返回最长的键；
否则从characters_list中寻找第一个匹配的文字，返回序列索引
获取指定坐标的单行文字
参数：
:param characters_list: 预选文字列表
:param points: 图像截取范围
Expand All @@ -632,7 +630,7 @@ def ocr_pos_sp(self, characters_list:list[str] = None, points = (0,0,0,0), numbe
:return index: 预选文字列表索引
:return data: 文字
"""
data = self.part_ocr_other(points, debug, number=number, img_pk=img_pk, is_single_line=True)
data = self.part_ocr(points, debug, number=number, img_pk=img_pk, is_single_line=True)
if len(data) == 0: # 字符串长度为零
return None
if not characters_list:
Expand All @@ -644,27 +642,6 @@ def ocr_pos_sp(self, characters_list:list[str] = None, points = (0,0,0,0), numbe
log.error(f"OCR失败，所识别的 '{data}' 不在指定预选项{characters_list}中")
return -1

def part_ocr(self,points = (0,0,0,0), debug=False, only_white=False):
"""
说明：
返回图片文字和坐标(相对于图片的坐标)
参数：
:param points: 图像截取范围
返回:
:return data: 文字: 坐标(相对于图片的坐标)
"""
img_fp, left, top, right, bottom, width, length = self.take_screenshot(points)
if only_white:
img_fp = self.remove_non_white_pixels(img_fp)
if debug:
show_img(img_fp)
#cv.imwrite("H://xqtd//xl//test.png",img_fp)
x, y = width/100*points[0], length/100*points[1]
out = self.ocr.ocr(img_fp)
data = {i['text'].replace(" ", ""): (int(left+x+(i['position'][2][0]+i['position'][0][0])/2),int(top+y+(i['position'][2][1]+i['position'][0][1])/2)) for i in out}
log.debug(data)
return data

def read_img(self, path, prefix='./picture/pc/'):
"""
说明：
Expand All @@ -676,7 +653,7 @@ def read_img(self, path, prefix='./picture/pc/'):
"""
return cv.imread(f'{prefix}{path}')

def part_ocr_other(self,points = (0,0,0,0), debug=False, left=False, number = False, img_pk:tuple = None, is_single_line = False
def part_ocr(self,points = (0,0,0,0), debug=False, left=False, number = False, img_pk:tuple = None, is_single_line = False, only_white=False
) -> Union[str, dict[str, tuple[int, int]]]:
"""
说明：
Expand All @@ -698,10 +675,8 @@ def part_ocr_other(self,points = (0,0,0,0), debug=False, left=False, number = Fa
int(width/100*points[0]):int(width/100*points[2]), :]
else:
img_fp, game_left, game_top, _, _, width, length = self.take_screenshot(points)
if debug:
# show_img(img_fp)
timestamp_str = str(int(datetime.timestamp(datetime.now())))
cv.imwrite(f"temp/relic_{str(points)}_{timestamp_str}.png", img_fp)
if only_white:
img_fp = self.remove_non_white_pixels(img_fp)
x, y = width/100*points[0], length/100*points[1]
if is_single_line:
out = self.number_ocr.ocr_for_single_line(img_fp) if number else self.ocr.ocr_for_single_line(img_fp)
Expand All @@ -714,6 +689,9 @@ def part_ocr_other(self,points = (0,0,0,0), debug=False, left=False, number = Fa
data = {i['text'].replace(" ", ""): (int(game_left+x+(i['position'][2][0]+i['position'][0][0])/2),int(game_top+y+(i['position'][2][1]+i['position'][0][1])/2)) for i in out}
if debug:
log.info(data)
# show_img(img_fp)
timestamp_str = str(int(datetime.timestamp(datetime.now())))
cv.imwrite(f"temp/relic_{str(points)}_{timestamp_str}.png", img_fp)
else:
log.debug(data)
return data
Expand Down
24 changes: 12 additions & 12 deletions utils/relic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ def equip_loadout(self, relics_hash:list[str]):
continue
# 点击装备
self.calculated.relative_click(pos)
button = self.calculated.ocr_pos_sp([_("装备"), _("替换"), _("卸下")], (75,90,82,95)) # 需识别[装备,替换,卸下]
button = self.calculated.ocr_pos_for_singleLine([_("装备"), _("替换"), _("卸下")], (75,90,82,95)) # 需识别[装备,替换,卸下]
if button in [0,1]:
log.info(_("点击装备"))
self.calculated.relative_click((78,92))
Expand Down Expand Up @@ -267,7 +267,7 @@ def save_loadout(self, character_name:str=None, max_retries=3):
raise Exception(_("重试次数达到上限"))
retry += 1
log.info(_(f"第 {retry} 次尝试重新OCR"))
tmp_hash = self.calculated.get_relic_hash(tmp_data)
tmp_hash = self.calculated.get_data_hash(tmp_data)
log.debug("\n"+pp.pformat(tmp_data))
self.print_relic(tmp_data)
if tmp_hash in self.relics_data:
Expand Down Expand Up @@ -343,7 +343,7 @@ def search_relic(self, equip_indx:int, key_hash:str=None, key_data:dict=None, ov
retry += 1
log.info(_(f"第 {retry} 次尝试重新OCR"))
# log.info("\n"+pp.pformat(tmp_data))
tmp_hash = self.calculated.get_relic_hash(tmp_data)
tmp_hash = self.calculated.get_data_hash(tmp_data)
if key_hash and key_hash == tmp_hash: # 精确搜索
return (x, y)
if key_data and self.compare_relics(key_data, tmp_data): # 模糊搜索
Expand Down Expand Up @@ -392,7 +392,7 @@ def add_relic_data(self, data:dict, data_hash:str=None) -> bool:
录入仪器数据
"""
if not data_hash:
data_hash = self.calculated.get_relic_hash(data)
data_hash = self.calculated.get_data_hash(data)
if data_hash not in self.relics_data:
self.relics_data = modify_json_file(RELIC_FILE_NAME, data_hash, data) # 返回更新后的字典
return True
Expand All @@ -407,7 +407,7 @@ def ocr_character_name(self) -> str:
返回：
:return character_name: 人物名称
"""
str = self.calculated.ocr_pos_sp(points=(13,4,22,9)) # 识别人物名称 (主角名称为玩家自定义，无法适用预选列表)
str = self.calculated.ocr_pos_for_singleLine(points=(13,4,22,9)) # 识别人物名称 (主角名称为玩家自定义，无法适用预选列表)
character_name = re.sub(r"[.’,，。、·'\"/\\]", '', str) # 删除由于背景光点造成的误判
log.info(_(f"识别人物: {character_name}"))
if character_name not in self.loadout_data:
Expand All @@ -430,25 +430,25 @@ def ocr_relic(self, equip_set_index:int = None) -> dict:
img_pc = self.calculated.take_screenshot() # 仅截取一次图片
# [1]部位识别
if equip_set_index is None:
equip_set_index = self.calculated.ocr_pos_sp(self.equip_set_name, points=(71,22,78,26), img_pk=img_pc)
equip_set_index = self.calculated.ocr_pos_for_singleLine(self.equip_set_name, points=(71,22,78,26), img_pk=img_pc)
if equip_set_index < 0:
raise RelicOCRException(_("遗器套装OCR错误"))
equip_set_name = self.equip_set_name[equip_set_index]
# [2]套装识别
name_list = self.relic_set_name[:, 0].tolist()
relic_set_index = self.calculated.ocr_pos_sp(name_list, points=(71,17,88,21), img_pk=img_pc)
relic_set_index = self.calculated.ocr_pos_for_singleLine(name_list, points=(71,17,88,21), img_pk=img_pc)
if relic_set_index < 0:
raise RelicOCRException(_("遗器部位OCR错误"))
relic_set_name = self.relic_set_name[relic_set_index, -1]
# [3]等级识别
level = self.calculated.ocr_pos_sp(points=(94,22,98,26), number=True, img_pk=img_pc)
level = self.calculated.ocr_pos_for_singleLine(points=(94,22,98,26), number=True, img_pk=img_pc)
level = int(level.split('+')[-1]) # 消除开头可能的'+'号
if level > 15:
raise RelicOCRException(_("遗器等级OCR错误"))
# [4]主属性识别
name_list = self.base_stats_name4equip[equip_set_index][:, 0].tolist()
base_stats_index = self.calculated.ocr_pos_sp(name_list, points=(74,29,89,34), img_pk=img_pc)
base_stats_value = self.calculated.ocr_pos_sp(points=(91,29,98,34), number=True, img_pk=img_pc)
base_stats_index = self.calculated.ocr_pos_for_singleLine(name_list, points=(74,29,89,34), img_pk=img_pc)
base_stats_value = self.calculated.ocr_pos_for_singleLine(points=(91,29,98,34), number=True, img_pk=img_pc)
if base_stats_index < 0:
raise RelicOCRException(_("遗器主词条OCR错误"))
if base_stats_value is None:
Expand All @@ -467,11 +467,11 @@ def ocr_relic(self, equip_set_index:int = None) -> dict:
subs_stats_dict = {}
total_level = 0
for name_point, value_point in zip(subs_stats_name_points, subs_stats_value_points):
tmp_index = self.calculated.ocr_pos_sp(name_list, points=name_point, img_pk=img_pc)
tmp_index = self.calculated.ocr_pos_for_singleLine(name_list, points=name_point, img_pk=img_pc)
if tmp_index is None: break # 所识别data为空，即词条为空，正常退出循环
if tmp_index < 0:
raise RelicOCRException(_("遗器副词条OCR错误"))
tmp_value = self.calculated.ocr_pos_sp(points=value_point, number=True, img_pk=img_pc)
tmp_value = self.calculated.ocr_pos_for_singleLine(points=value_point, number=True, img_pk=img_pc)
if tmp_value is None:
raise RelicOCRException(_("遗器副词条数值OCR错误"))
else:
Expand Down

0 comments on commit e50983d

Please sign in to comment.