Add image processing with threshold and contours

main.py

@@ -1,23 +1,52 @@
+from PIL import Image
 from pyboy import PyBoy
-from pyboy.utils import WindowEvent
 
-valid_buttons = ["a", "b", "start", "select", "left", "right", "up", "down"]
+from utils.game_elements import detect_game_elements
+from utils.image_processing import detect_contours, process_image
 
 
 def load_state(path: str, pyboy: PyBoy) -> None:
+    """
+    Load the state of the game from a file.
+
+    Args:
+        path (str): The path to the file containing the game state.
+        pyboy (PyBoy): The PyBoy instance representing the game.
+
+    Returns:
+        None
+    """
     with open(path, "rb") as state:
         pyboy.load_state(state)
 
 
+def get_game_elemets(screen_image: Image.Image) -> dict:
+    """
+    Extracts game elements from a screen image.
+
+    Args:
+        screen_image (PIL.Image.Image): The screen image to process.
+
+    Returns:
+        dict: A dictionary containing the detected game elements.
+    """
+    thresh_image, screen_np = process_image(screen_image)
+    contours, contour_image = detect_contours(thresh_image, screen_np)
+    game_elements = detect_game_elements(contours)
+
+    return game_elements
+
+
 if __name__ == "__main__":
     pyboy = PyBoy("roms/ZeldaLinksAwakening.gb")
 
     load_state("roms/ZeldaLinksAwakening.gb.state", pyboy)
 
     while pyboy.tick():
-        pyboy.button(valid_buttons[0])  # a
         pyboy.tick()
-        pyboy.screen.image.save("captures/test.png")
+        screen_image = pyboy.screen.image
+        game_elements = get_game_elemets(screen_image)
+        print(game_elements)
         pass
 
     pyboy.stop()

utils/game_elements.py

@@ -0,0 +1,47 @@
+import cv2
+
+
+def detect_game_elements(contours: tuple) -> dict:
+    """
+    Detects game elements from a list of contours.
+
+    Args:
+        contours (tuple): A tuple of contours.
+
+    Returns:
+        dict: A dictionary containing the detected game elements.
+            The dictionary has the following structure:
+            {
+                "enemies": [(x, y, w, h), ...],
+                "walls": [(x, y, w, h), ...],
+                "link": (x, y, w, h)
+            }
+            - "enemies" is a list of enemy bounding boxes.
+            - "walls" is a list of wall bounding boxes.
+            - "link" is the bounding box of the player character.
+    """
+
+    def is_enemy(x, y, w, h):
+        return w > 15 and h > 15
+
+    def is_wall(x, y, w, h):
+        return w > 18 and h > 18
+
+    def is_link(x, y, w, h):
+        return w > 12 and h > 12
+
+    game_elements = {"enemies": [], "walls": [], "link": None}
+
+    for contour in contours:
+        area = cv2.contourArea(contour)
+        if area > 50:
+            x, y, w, h = cv2.boundingRect(contour)
+            if w > 10 and h > 10:
+                if is_enemy(x, y, w, h):
+                    game_elements["enemies"].append((x, y, w, h))
+                elif is_wall(x, y, w, h):
+                    game_elements["walls"].append((x, y, w, h))
+                elif is_link(x, y, w, h):
+                    game_elements["link"] = (x, y, w, h)
+
+    return game_elements

utils/image_processing.py

@@ -0,0 +1,50 @@
+from typing import Union
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def process_image(screen_image: Image.Image) -> Union[Image.Image, np.array]:
+    """
+    Process the given screen image.
+
+    Args:
+        screen_image (PIL.Image.Image): The input screen image.
+
+    Returns:
+        Union[PIL.Image.Image, np.array]: The processed image as a PIL Image and a NumPy array.
+    """
+    screen_np = np.array(screen_image)
+
+    gray_image = cv2.cvtColor(screen_np, cv2.COLOR_BGR2GRAY)
+
+    _, thresh_image = cv2.threshold(gray_image, 127, 255, cv2.THRESH_BINARY)
+
+    screen_np = np.array(screen_np)
+
+    cv2.imwrite("captures/thresh.png", thresh_image)
+
+    return (thresh_image, screen_np)
+
+
+def detect_contours(
+    thresh_image: Image.Image, screen_np: np.array
+) -> Union[tuple, np.ndarray]:
+    """
+    Detects contours in a thresholded image and draws them on the original image.
+
+    Args:
+        thresh_image (PIL.Image.Image): The thresholded image.
+        screen_np (numpy.ndarray): The original image as a numpy array.
+
+    Returns:
+        Union[tuple, numpy.ndarray]: A tuple containing the detected contours and the image with contours drawn on it.
+    """
+    contours, _ = cv2.findContours(thresh_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+    contour_image = cv2.drawContours(screen_np.copy(), contours, -1, (0, 255, 0), 2)
+
+    cv2.imwrite("captures/contour.png", contour_image)
+
+    return (contours, contour_image)