From 35a5d63034f88032206392c45f2fce3ca2d6b0a4 Mon Sep 17 00:00:00 2001 From: Nikolas Date: Mon, 23 Sep 2024 14:59:05 -0700 Subject: [PATCH] initial commit, linter come at me --- CHANGELOG.md | 4 + F1_22/f1.py | 4 +- F1_23/f1_23.py | 5 +- atomic_heart/atomicheart.py | 3 +- cities_skylines_2/citiesskylines2.py | 2 +- counterstrike2/cs2.py | 3 +- cyberpunk2077/cyberpunk2077.py | 3 +- dota2/dota2.py | 3 +- forza5/forza5.py | 3 +- grid_legends/grid_legends.py | 2 +- harness_utils/keras_service.py | 155 +++++++++++++++--- returnal/returnal.py | 3 +- rocket_league/rocket_league.py | 3 +- .../shadowofthetombraider.py | 2 +- stellaris/stellaris.py | 2 +- .../the_last_of_us_part_i.py | 3 +- total_war_pharaoh/total_war_pharaoh/twp.py | 2 +- total_war_warhammer_iii/twwh3.py | 3 +- 18 files changed, 151 insertions(+), 54 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bb910ca..73432c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. Changes are grouped by the date they are merged to the main branch of the repository and are ordered from newest to oldest. Dates use the ISO 8601 extended calendar date format, i.e. YYYY-MM-DD. +## 2024-09-23 + +- Add screen splitting to Keras Service. + ## 2024-9-20 - Add Grid Legends test harness. diff --git a/F1_22/f1.py b/F1_22/f1.py index 7f2d0fe..3ce77ef 100644 --- a/F1_22/f1.py +++ b/F1_22/f1.py @@ -158,9 +158,7 @@ console.setFormatter(formatter) logging.getLogger("").addHandler(console) args = get_args() -kerasService = KerasService( - args.keras_host, args.keras_port, os.path.join(LOG_DIRECTORY, "screenshot.jpg") -) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/F1_23/f1_23.py b/F1_23/f1_23.py index e4e19b4..3cb0b6e 100644 --- a/F1_23/f1_23.py +++ b/F1_23/f1_23.py @@ -228,10 +228,7 @@ parser.add_argument( "--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True ) args = parser.parse_args() -kerasService = KerasService( - args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg") -) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/atomic_heart/atomicheart.py b/atomic_heart/atomicheart.py index bf13aba..1f446a4 100644 --- a/atomic_heart/atomicheart.py +++ b/atomic_heart/atomicheart.py @@ -141,8 +141,7 @@ parser.add_argument("--kerasHost", dest="keras_host", parser.add_argument("--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True) args = parser.parse_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/cities_skylines_2/citiesskylines2.py b/cities_skylines_2/citiesskylines2.py index 324c9fb..00c29f5 100644 --- a/cities_skylines_2/citiesskylines2.py +++ b/cities_skylines_2/citiesskylines2.py @@ -123,7 +123,7 @@ def main(): parser.add_argument("--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True) args = parser.parse_args() - keras_service = KerasService(args.keras_host, args.keras_port, LOG_DIR.joinpath("screenshot.jpg")) + keras_service = KerasService(args.keras_host, args.keras_port) test_start_time, test_end_time = run_benchmark(keras_service) resolution = read_current_resolution() diff --git a/counterstrike2/cs2.py b/counterstrike2/cs2.py index b9df156..f8829bd 100644 --- a/counterstrike2/cs2.py +++ b/counterstrike2/cs2.py @@ -125,8 +125,7 @@ def main(): help="Port for Keras OCR service", required=True) args = parser.parse_args() - keras_service = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIR, "screenshot.jpg")) + keras_service = KerasService(args.keras_host, args.keras_port) start_time, end_time = run_benchmark(keras_service) diff --git a/cyberpunk2077/cyberpunk2077.py b/cyberpunk2077/cyberpunk2077.py index e41580d..2652bd6 100644 --- a/cyberpunk2077/cyberpunk2077.py +++ b/cyberpunk2077/cyberpunk2077.py @@ -126,8 +126,7 @@ console.setFormatter(formatter) logging.getLogger('').addHandler(console) args = get_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/dota2/dota2.py b/dota2/dota2.py index e464634..7bda0ed 100644 --- a/dota2/dota2.py +++ b/dota2/dota2.py @@ -37,8 +37,7 @@ console.setFormatter(formatter) logging.getLogger('').addHandler(console) args = get_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) def start_game(): diff --git a/forza5/forza5.py b/forza5/forza5.py index d4d2465..4c21d26 100644 --- a/forza5/forza5.py +++ b/forza5/forza5.py @@ -133,8 +133,7 @@ parser.add_argument("--kerasHost", dest="keras_host", parser.add_argument("--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True) args = parser.parse_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/grid_legends/grid_legends.py b/grid_legends/grid_legends.py index e451f6b..3fbde15 100644 --- a/grid_legends/grid_legends.py +++ b/grid_legends/grid_legends.py @@ -142,7 +142,7 @@ def run_benchmark(keras_service): def main(): """entry point""" args = get_args() - keras_service = KerasService(args.keras_host, args.keras_port, LOG_DIR.joinpath("screenshot.jpg")) + keras_service = KerasService(args.keras_host, args.keras_port) start_time, end_time = run_benchmark(keras_service) elapsed_test_time = round((end_time - start_time), 2) logging.info("Benchmark took %f seconds", elapsed_test_time) diff --git a/harness_utils/keras_service.py b/harness_utils/keras_service.py index 7ac5032..e5b18ac 100644 --- a/harness_utils/keras_service.py +++ b/harness_utils/keras_service.py @@ -1,45 +1,90 @@ """Allows accessing Keras Service if available.""" +import io import json import logging -import os import time import mss import cv2 import requests import numpy as np +from enum import Enum +from dataclasses import dataclass logging.getLogger("requests").setLevel(logging.WARNING) logging.getLogger("urllib3").setLevel(logging.WARNING) DEFAULT_TIMEOUT = 120.0 + +class ScreenShotDivideMethod(Enum): + """split method""" + horizontal = "horizontal" + vertical = "vertical" + quadrant = "quadrant" + none = "none" + + +class ScreenShotQuadrant(Enum): + """split arguments""" + TOP = 1 + BOTTOM = 2 + LEFT = 1 + RIGHT = 2 + TOP_LEFT = 1 + TOP_RIGHT = 2 + BOTTOM_LEFT = 3 + BOTTOM_RIGHT = 4 + + +class FrameDivideException(ValueError): + """Exception indicating was encountered while trying to divide a frame""" + + +@dataclass +class ScreenSplitConfig: + """data class to contain config for taking splitting a screen shot""" + divide_method: ScreenShotDivideMethod + quadrant: ScreenShotQuadrant + + class KerasService(): """Sets up connection to a Keras service and provides methods to use it""" + def __init__( self, ip_addr: str, port: int | str, - screenshot_path: str | os.PathLike, timeout: float = DEFAULT_TIMEOUT) -> None: self.ip_addr = ip_addr self.port = str(port) self.url = f"http://{ip_addr}:{str(port)}/process" - self.screenshot_path = screenshot_path self.timeout = timeout - def _capture_screenshot(self) -> None: + def _capture_screenshot(self, split_config: ScreenSplitConfig) -> io.BytesIO: + screenshot = None with mss.mss() as sct: monitor_1 = sct.monitors[1] # Identify the display to capture - screen = np.array(sct.grab(monitor_1)) - screen = cv2.cvtColor(screen, cv2.COLOR_RGB2GRAY) - cv2.imwrite(str(self.screenshot_path), screen) + screenshot = np.array(sct.grab(monitor_1)) + screenshot = cv2.cvtColor(screenshot, cv2.COLOR_RGB2GRAY) - def _query_service(self, word: str, report_file: any) -> any: + if split_config.divide_method == ScreenShotDivideMethod.horizontal: + screenshot = self._divide_horizontal(screenshot) + elif split_config.divide_method == ScreenShotDivideMethod.vertical: + screenshot = self._divide_vertical(screenshot) + elif split_config.divide_method == ScreenShotDivideMethod.quadrant: + screenshot = self._divide_in_four(screenshot) + elif split_config.divide_method == ScreenShotDivideMethod.none: + screenshot = screenshot + + _, encoded_image = cv2.imencode('.jpg', screenshot) + return io.BytesIO(encoded_image) + + def _query_service(self, word: str, image_bytes: io.BytesIO) -> any: try: keras_response = requests.post( self.url, data={"word": word}, - files={"file": report_file}, + files={"file": image_bytes}, timeout=self.timeout ) @@ -53,11 +98,56 @@ class KerasService(): except requests.exceptions.Timeout: return None - def capture_screenshot_find_word(self, word: str) -> any: - """Take a screenshot and try to find the given word within it.""" - self._capture_screenshot() - with open(self.screenshot_path, "rb") as report_file: - return self._query_service(word, report_file) + def _divide_horizontal(self, screenshot, quadrant: ScreenShotQuadrant): + """divide the screenshot horizontally""" + height, _, _ = screenshot.shape + if quadrant == ScreenShotQuadrant.TOP: + return screenshot[0:int(height/2), :] + elif quadrant == ScreenShotQuadrant.BOTTOM: + return screenshot[int(height/2):int(height), :] + else: + raise FrameDivideException( + f"Unrecognized quadrant for horizontal: {self.quadrant}") + + def _divide_vertical(self, screenshot, quadrant: ScreenShotQuadrant): + """divide the screenshot vertically""" + _, width, _ = screenshot.shape + if self.quadrant == quadrant.LEFT: + return screenshot[:, 0:int(width/2)] + elif self.quadrant == quadrant.RIGHT: + return screenshot[:, int(width/2):int(width)] + else: + raise FrameDivideException( + f"Unrecognized quadrant for vertical: {self.quadrant}") + + def _divide_in_four(self, screenshot, quadrant: ScreenShotQuadrant): + """divide the screenshot in four quadrants""" + height, width, _ = screenshot.shape + if self.quadrant == ScreenShotQuadrant.TOP_LEFT: + return screenshot[0:int(height/2), 0:int(width/2)] + elif self.quadrant == ScreenShotQuadrant.TOP_RIGHT: + return screenshot[0:int(height/2), int(width/2):int(width)] + elif self.quadrant == ScreenShotQuadrant.BOTTOM_LEFT: + return screenshot[int(height/2):int(height), 0:int(width/2)] + elif self.quadrant == ScreenShotQuadrant.BOTTOM_RIGHT: + return screenshot[int(height/2):int(height), int(width/2):int(width)] + else: + raise FrameDivideException( + f"Unrecognized quadrant for in four: {self.quadrant}") + + def _look_for_word( + self, word: str, + attempts: int = 1, + interval: float = 0.0, + split_config: ScreenSplitConfig = None) -> bool: + """implementation of look for word""" + for _ in range(attempts): + image_bytes = self._capture_screenshot(split_config) + result = self._query_service(word, image_bytes) + if result is not None: + return result + time.sleep(interval) + return None def look_for_word(self, word: str, attempts: int = 1, interval: float = 0.0) -> bool: """Takes a screenshot of the monitor and searches for a given word. @@ -65,23 +155,40 @@ class KerasService(): of attempts is reached. Will return early if the query result comes back with a match. """ - for _ in range(attempts): - result = self.capture_screenshot_find_word(word) + return self._look_for_word(word, attempts, interval, ScreenSplitConfig(divide_method=ScreenShotDivideMethod.none)) + + def look_for_word(self, word: str, attempts: int = 1, interval: float = 0.0, split_config: ScreenSplitConfig = None) -> bool: + """Overload for look_for_word but allows for screen splitting + which will look for a word in only part of the screen + """ + return self._look_for_word(word, attempts, interval, split_config) + + def _wait_for_word( + self, + word: str, + interval: float, + timeout: float, + split_config: ScreenSplitConfig) -> bool: + """implementation of wait for word""" + search_start_time = time.time() + while time.time() - search_start_time < timeout: + image_bytes = self._capture_screenshot(split_config) + result = self._query_service(word, image_bytes) if result is not None: return result time.sleep(interval) return None - def wait_for_word(self, word: str, interval: float = 0.0, timeout: float = 0.0): + def wait_for_word(self, word: str, interval: float = 0.0, timeout: float = 0.0) -> bool: """Takes a screenshot of the monitor and searches for a given word. Will look for the word at a given time interval until the specified timeout has been exceeded. Will return early if the query result comes back with a match. """ - search_start_time = time.time() - while time.time() - search_start_time < timeout: - result = self.capture_screenshot_find_word(word) - if result is not None: - return result - time.sleep(interval) - return None + return self._wait_for_word(word, interval, timeout, ScreenSplitConfig(divide_method=ScreenShotDivideMethod.none)) + + def wait_for_word(self, word: str, interval: float = 0.0, timeout: float = 0.0, split_config: ScreenSplitConfig = None) -> bool: + """Overload for wait_for_word but allows for screen splitting + which will look for a word in only part of the screen + """ + return self._wait_for_word(word, interval, timeout, split_config) diff --git a/returnal/returnal.py b/returnal/returnal.py index 26030f5..ed7b3c4 100644 --- a/returnal/returnal.py +++ b/returnal/returnal.py @@ -156,8 +156,7 @@ console.setFormatter(formatter) logging.getLogger('').addHandler(console) args = get_args() -kerasService = KerasService( - args.keras_host, args.keras_port, os.path.join(LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/rocket_league/rocket_league.py b/rocket_league/rocket_league.py index a929c6c..f9dd4da 100644 --- a/rocket_league/rocket_league.py +++ b/rocket_league/rocket_league.py @@ -38,8 +38,7 @@ console.setFormatter(formatter) logging.getLogger('').addHandler(console) args = get_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) def get_run_game_id_command(game_id: int) -> str: diff --git a/shadowofthetombraider/shadowofthetombraider.py b/shadowofthetombraider/shadowofthetombraider.py index eb967f3..d72ff94 100644 --- a/shadowofthetombraider/shadowofthetombraider.py +++ b/shadowofthetombraider/shadowofthetombraider.py @@ -52,7 +52,7 @@ def run_benchmark(): start_game() args = get_args() - keras_service = KerasService(args.keras_host, args.keras_port, LOG_DIR.joinpath("screenshot.jpg")) + keras_service = KerasService(args.keras_host, args.keras_port) am = ArtifactManager(LOG_DIR) if keras_service.wait_for_word(word="options", timeout=30, interval=1) is None: diff --git a/stellaris/stellaris.py b/stellaris/stellaris.py index 120cc23..ba9ac0e 100644 --- a/stellaris/stellaris.py +++ b/stellaris/stellaris.py @@ -61,7 +61,7 @@ def console_command(command): def run_benchmark(keras_host, keras_port): """Starts the benchmark""" - keras_service = KerasService(keras_host, keras_port, str(LOG_DIR.joinpath("screenshot.jpg"))) + keras_service = KerasService(keras_host, keras_port) copy_benchmarkfiles() copy_benchmarksave() start_game() diff --git a/the_last_of_us_part_i/the_last_of_us_part_i.py b/the_last_of_us_part_i/the_last_of_us_part_i.py index f930eb9..130b437 100644 --- a/the_last_of_us_part_i/the_last_of_us_part_i.py +++ b/the_last_of_us_part_i/the_last_of_us_part_i.py @@ -126,8 +126,7 @@ console.setFormatter(formatter) logging.getLogger('').addHandler(console) args = get_args() -kerasService = KerasService( - args.keras_host, args.keras_port, os.path.join(LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, end_time = run_benchmark() diff --git a/total_war_pharaoh/total_war_pharaoh/twp.py b/total_war_pharaoh/total_war_pharaoh/twp.py index 4f9ed7f..e57407d 100644 --- a/total_war_pharaoh/total_war_pharaoh/twp.py +++ b/total_war_pharaoh/total_war_pharaoh/twp.py @@ -167,7 +167,7 @@ def main(): parser.add_argument("--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True) args = parser.parse_args() - keras_service = KerasService(args.keras_host, args.keras_port, LOG_DIR.joinpath("screenshot.jpg")) + keras_service = KerasService(args.keras_host, args.keras_port) start_time, endtime = run_benchmark(keras_service) height, width = read_current_resolution() report = { diff --git a/total_war_warhammer_iii/twwh3.py b/total_war_warhammer_iii/twwh3.py index db72241..0a38723 100644 --- a/total_war_warhammer_iii/twwh3.py +++ b/total_war_warhammer_iii/twwh3.py @@ -172,8 +172,7 @@ parser.add_argument("--kerasHost", dest="keras_host", parser.add_argument("--kerasPort", dest="keras_port", help="Port for Keras OCR service", required=True) args = parser.parse_args() -kerasService = KerasService(args.keras_host, args.keras_port, os.path.join( - LOG_DIRECTORY, "screenshot.jpg")) +kerasService = KerasService(args.keras_host, args.keras_port) try: start_time, endtime = run_benchmark()