diff --git a/assets/share/planner/result/CALCULATE_TITLE.png b/assets/share/planner/result/CALCULATE_TITLE.png new file mode 100644 index 000000000..5850f57a1 Binary files /dev/null and b/assets/share/planner/result/CALCULATE_TITLE.png differ diff --git a/assets/share/planner/result/DETAIL_TITLE.png b/assets/share/planner/result/DETAIL_TITLE.png new file mode 100644 index 000000000..f4386bf44 Binary files /dev/null and b/assets/share/planner/result/DETAIL_TITLE.png differ diff --git a/assets/share/planner/result/MATERIAL_TITLE.png b/assets/share/planner/result/MATERIAL_TITLE.png new file mode 100644 index 000000000..a60a04717 Binary files /dev/null and b/assets/share/planner/result/MATERIAL_TITLE.png differ diff --git a/assets/share/planner/result/OCR_RESULT.png b/assets/share/planner/result/OCR_RESULT.png new file mode 100644 index 000000000..14f8d7fe3 Binary files /dev/null and b/assets/share/planner/result/OCR_RESULT.png differ diff --git a/assets/share/planner/result/RESULT_CHECK.SEARCH.png b/assets/share/planner/result/RESULT_CHECK.SEARCH.png new file mode 100644 index 000000000..49da0570b Binary files /dev/null and b/assets/share/planner/result/RESULT_CHECK.SEARCH.png differ diff --git a/assets/share/planner/result/RESULT_CHECK.png b/assets/share/planner/result/RESULT_CHECK.png new file mode 100644 index 000000000..6b183abbd Binary files /dev/null and b/assets/share/planner/result/RESULT_CHECK.png differ diff --git a/assets/share/planner/result/RESULT_SCROLL.png b/assets/share/planner/result/RESULT_SCROLL.png new file mode 100644 index 000000000..5b8e7a042 Binary files /dev/null and b/assets/share/planner/result/RESULT_SCROLL.png differ diff --git a/tasks/planner/assets/assets_planner_result.py b/tasks/planner/assets/assets_planner_result.py new file mode 100644 index 000000000..46dc97ab0 --- /dev/null +++ b/tasks/planner/assets/assets_planner_result.py @@ -0,0 +1,65 @@ +from module.base.button import Button, ButtonWrapper + +# This file was auto-generated, do not modify it manually. To generate: +# ``` python -m dev_tools.button_extract ``` + +CALCULATE_TITLE = ButtonWrapper( + name='CALCULATE_TITLE', + share=Button( + file='./assets/share/planner/result/CALCULATE_TITLE.png', + area=(264, 187, 346, 207), + search=(244, 167, 366, 227), + color=(141, 143, 161), + button=(264, 187, 346, 207), + ), +) +DETAIL_TITLE = ButtonWrapper( + name='DETAIL_TITLE', + share=Button( + file='./assets/share/planner/result/DETAIL_TITLE.png', + area=(263, 402, 346, 422), + search=(243, 382, 366, 442), + color=(144, 146, 164), + button=(263, 402, 346, 422), + ), +) +MATERIAL_TITLE = ButtonWrapper( + name='MATERIAL_TITLE', + share=Button( + file='./assets/share/planner/result/MATERIAL_TITLE.png', + area=(263, 194, 346, 214), + search=(243, 174, 366, 234), + color=(135, 136, 156), + button=(263, 194, 346, 214), + ), +) +OCR_RESULT = ButtonWrapper( + name='OCR_RESULT', + share=Button( + file='./assets/share/planner/result/OCR_RESULT.png', + area=(331, 104, 1201, 592), + search=(311, 84, 1221, 612), + color=(254, 254, 254), + button=(331, 104, 1201, 592), + ), +) +RESULT_CHECK = ButtonWrapper( + name='RESULT_CHECK', + share=Button( + file='./assets/share/planner/result/RESULT_CHECK.png', + area=(263, 128, 324, 148), + search=(255, 103, 355, 593), + color=(132, 119, 92), + button=(263, 128, 324, 148), + ), +) +RESULT_SCROLL = ButtonWrapper( + name='RESULT_SCROLL', + share=Button( + file='./assets/share/planner/result/RESULT_SCROLL.png', + area=(1238, 104, 1245, 592), + search=(1218, 84, 1265, 612), + color=(33, 44, 62), + button=(1238, 104, 1245, 592), + ), +) diff --git a/tasks/planner/result.py b/tasks/planner/result.py new file mode 100644 index 000000000..e4725e3aa --- /dev/null +++ b/tasks/planner/result.py @@ -0,0 +1,202 @@ +import cv2 +from pponnxcr.predict_system import BoxedResult +from pydantic import BaseModel + +from module.base.utils import area_center, area_in_area +from module.logger import logger +from module.ocr.ocr import OcrWhiteLetterOnComplexBackground +from module.ui.scroll import AdaptiveScroll +from tasks.daily.synthesize import SynthesizeUI +from tasks.planner.assets.assets_planner_result import * +from tasks.planner.keywords import ITEM_CLASSES +from tasks.planner.keywords.classes import ItemBase, ItemCurrency + +CALCULATE_TITLE.load_search(RESULT_CHECK.search) +MATERIAL_TITLE.load_search(RESULT_CHECK.search) +DETAIL_TITLE.load_search(RESULT_CHECK.search) + + +class PlannerResultRow(BaseModel): + item: ItemBase + total: int + synthesize: int + demand: int + + def __eq__(self, other): + return self.item == other.item + + +class OcrPlannerResult(OcrWhiteLetterOnComplexBackground): + def __init__(self): + # Planner currently CN only + super().__init__(OCR_RESULT, lang='cn') + self.limited_area = OCR_RESULT.area + self.limit_y = 720 + + def filter_detected(self, result: BoxedResult) -> bool: + if not area_in_area(result.box, self.limited_area, threshold=0): + return False + if area_center(result.box)[1] > self.limit_y: + return False + return True + + def detect_and_ocr(self, image, *args, **kwargs): + # Remove rows below DETAIL_TITLE + if DETAIL_TITLE.match_template(image): + self.limit_y = DETAIL_TITLE.button[3] + else: + self.limit_y = 720 + return super().detect_and_ocr(image, *args, **kwargs) + + def pre_process(self, image): + r, g, b = cv2.split(image) + cv2.max(r, g, dst=r) + cv2.max(r, b, dst=r) + image = cv2.merge([r, r, r]) + return image + + +class PlannerResult(SynthesizeUI): + def is_in_planner_result(self): + if self.appear(RESULT_CHECK): + return True + if self.appear(CALCULATE_TITLE): + return True + if self.appear(MATERIAL_TITLE): + return True + if self.appear(DETAIL_TITLE): + return True + return False + + def parse_planner_result_page(self) -> list[PlannerResultRow]: + """ + Pages: + in: planner result + """ + ocr = OcrPlannerResult() + results = ocr.detect_and_ocr(self.device.image) + + x_total = 842 + x_synthesize = 965 + x_demand = 1129 + + def x_match(result: BoxedResult, x): + rx = area_center(result.box)[0] + return x - 50 <= rx <= x + 50 + + def y_match(result: BoxedResult, y): + rx = area_center(result.box)[1] + return y - 15 <= rx <= y + 15 + + # Split columns + list_item = [r for r in results + if not r.ocr_text.isdigit() and ocr._match_result(r.ocr_text, keyword_classes=ITEM_CLASSES)] + list_number = [r for r in results if r.ocr_text.isdigit()] + list_total = [r for r in list_number if x_match(r, x_total)] + list_synthesize = [r for r in list_number if x_match(r, x_synthesize)] + list_demand = [r for r in list_number if x_match(r, x_demand)] + + # Structure + out: list[PlannerResultRow] = [] + for item in list_item: + y_item = area_center(item.box)[1] + total = -1 + for number in list_total: + if y_match(number, y_item): + total = int(number.ocr_text) + break + synthesize = -1 + for number in list_synthesize: + if y_match(number, y_item): + synthesize = int(number.ocr_text) + break + demand = -1 + for number in list_demand: + if y_match(number, y_item): + demand = int(number.ocr_text) + break + item = ocr._match_result(item.ocr_text, keyword_classes=ITEM_CLASSES) + row = PlannerResultRow( + item=item, + total=total, + synthesize=synthesize, + demand=demand + ) + # Validate item + # print(row) + if row.total <= 0: + logger.warning(f'Planner row with total <= 0, {row}') + continue + if row.synthesize < 0: + # Credits always have synthesize="-" + if row.item.__class__ != ItemCurrency: + logger.warning(f'Planner row with synthesize < 0, {row}') + continue + if row.demand < 0: + logger.warning(f'Planner row with demand < 0, {row}') + continue + # Add + out.append(row) + + logger.info(f'parse_planner_result_page: {out}') + return out + + def parse_planner_result(self, skip_first_screenshot=True) -> list[PlannerResultRow]: + """ + Pages: + in: planner result + """ + logger.hr('Parse planner result') + scroll = AdaptiveScroll(RESULT_SCROLL.button, name=RESULT_SCROLL.name) + scroll.drag_threshold = 0.1 + scroll.edge_threshold = 0.1 + scroll.parameters = { + 'height': 50, + 'prominence': 15, + 'width': 5, + } + if not skip_first_screenshot: + self.device.screenshot() + skip_first_screenshot = False + if not scroll.at_top(main=self): + scroll.set_top(main=self) + + out = [] + while 1: + if skip_first_screenshot: + skip_first_screenshot = False + else: + self.device.screenshot() + + # Skip first page + if self.appear(CALCULATE_TITLE): + scroll.next_page(main=self, page=0.75) + continue + + # Parse + rows = self.parse_planner_result_page() + for row in rows: + if row not in out: + out.append(row) + logger.attr('PlannerResult', len(rows)) + + # Scroll + if scroll.at_bottom(main=self): + logger.info('Reached scroll end, stop') + break + elif self.appear(DETAIL_TITLE): + logger.info('Reached DETAIL_TITLE, stop') + break + else: + scroll.next_page(main=self, page=0.8) + + logger.hr('Planner Result') + for row in out: + logger.info(f'Item: {row.item.name}, {row.total}, {row.synthesize}, {row.demand}') + return out + + +if __name__ == '__main__': + self = PlannerResult('src') + self.device.screenshot() + self.parse_planner_result()