From 04c8c5eb7202aae4c16b9c3ff250b4e4a4d23e68 Mon Sep 17 00:00:00 2001
From: dffdeeq <dffdeeq@gmail.com>
Date: Wed, 24 Dec 2025 20:57:58 +0400
Subject: [PATCH] fix: default capture resolution to game window size (refs
 #10) and bbox format fix

---
 README.md            |  5 +++++
 nitrogen/game_env.py | 40 +++++++++++++++++++++++++---------------
 scripts/play.py      |  4 ++++
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 90cb439..8498f38 100644
--- a/README.md
+++ b/README.md
@@ -50,6 +50,11 @@ python scripts/play.py --process '<game_executable_name>.exe'
 
 The `--process` parameter must be the exact executable name of the game you want to play. You can find it by right-clicking on the game process in Windows Task Manager (Ctrl+Shift+Esc), and selecting `Properties`. The process name should be in the `General` tab and end with `.exe`.
 
+By default, screen capture uses the game window resolution. You can override it with `--width` and `--height`:
+```bash
+python scripts/play.py --process '<game_executable_name>.exe' --width 1920 --height 1080
+```
+
 <!-- TODO # Paper and Citation
 
 If you find our work useful, please consider citing us!
diff --git a/nitrogen/game_env.py b/nitrogen/game_env.py
index 982c1d7..f0557f0 100644
--- a/nitrogen/game_env.py
+++ b/nitrogen/game_env.py
@@ -350,12 +350,14 @@ def screenshot(self):
 
 
 class DxcamScreenshotBackend:
-    def __init__(self, bbox, fps):
+    def __init__(self, bbox_ltrb, fps):
         import dxcam
         self.camera = dxcam.create()
-        self.bbox = bbox
+        self.bbox_ltrb = bbox_ltrb
+        self.width = bbox_ltrb[2] - bbox_ltrb[0]
+        self.height = bbox_ltrb[3] - bbox_ltrb[1]
         self.last_screenshot = None
-        self.camera.start(region=self.bbox, target_fps=fps, video_mode=True)
+        self.camera.start(region=self.bbox_ltrb, target_fps=fps, video_mode=True)
 
     def screenshot(self):
         screenshot = self.camera.get_latest_frame()
@@ -364,7 +366,7 @@ def screenshot(self):
             if self.last_screenshot is not None:
                 return self.last_screenshot
             else:
-                return Image.new("RGB", (self.bbox[2], self.bbox[3]), (0, 0, 0))
+                return Image.new("RGB", (self.width, self.height), (0, 0, 0))
         screenshot = Image.fromarray(screenshot)
         self.last_screenshot = screenshot
         return screenshot
@@ -387,8 +389,8 @@ class GamepadEnv(Env):
     def __init__(
             self,
             game,
-            image_height=1440,
-            image_width=2560,
+            image_height=None,
+            image_width=None,
             controller_type="xbox",
             game_speed=1.0,
             env_fps=10,
@@ -404,8 +406,8 @@ def __init__(
         assert screenshot_backend in ["pyautogui", "dxcam"], "Screenshot backend must be either 'pyautogui' or 'dxcam'"
 
         self.game = game
-        self.image_height = int(image_height)
-        self.image_width = int(image_width)
+        self._image_height_override = image_height
+        self._image_width_override = image_width
         self.game_speed = game_speed
         self.env_fps = env_fps
         self.step_duration = self.calculate_step_duration()
@@ -424,10 +426,6 @@ def __init__(
         if self.game_pid is None:
             raise Exception(f"Could not find PID for game: {game}")
 
-        self.observation_space = Box(
-            low=0, high=255, shape=(self.image_height, self.image_width, 3), dtype="uint8"
-        )
-
         # Define a unified action space
         self.action_space = Dict(
             {
@@ -468,16 +466,28 @@ def __init__(
 
         self.game_window.activate()
         l, t, r, b = self.game_window.left, self.game_window.top, self.game_window.right, self.game_window.bottom
-        self.bbox = (l, t, r - l, b - t)
+        window_width, window_height = r - l, b - t
+
+        # Use window size if not explicitly set
+        self.image_width = self._image_width_override if self._image_width_override else window_width
+        self.image_height = self._image_height_override if self._image_height_override else window_height
+
+        # bbox formats: dxcam needs (l,t,r,b), pyautogui needs (l,t,w,h)
+        self.bbox_ltrb = (l, t, r, b)
+        self.bbox_ltwh = (l, t, window_width, window_height)
+
+        self.observation_space = Box(
+            low=0, high=255, shape=(self.image_height, self.image_width, 3), dtype="uint8"
+        )
 
         # Initialize speedhack client if using DLL injection
         self.speedhack_client = xsh.Client(process_id=self.game_pid, arch=self.game_arch)
 
         # Get the screenshot backend
         if screenshot_backend == "dxcam":
-            self.screenshot_backend = DxcamScreenshotBackend(self.bbox, self.env_fps)
+            self.screenshot_backend = DxcamScreenshotBackend(self.bbox_ltrb, self.env_fps)
         elif screenshot_backend == "pyautogui":
-            self.screenshot_backend = PyautoguiScreenshotBackend(self.bbox)
+            self.screenshot_backend = PyautoguiScreenshotBackend(self.bbox_ltwh)
         else:
             raise ValueError("Unsupported screenshot backend. Use 'dxcam' or 'pyautogui'.")
 
diff --git a/scripts/play.py b/scripts/play.py
index c0db5f9..72e959c 100644
--- a/scripts/play.py
+++ b/scripts/play.py
@@ -19,6 +19,8 @@
 parser.add_argument("--process", type=str, default="celeste.exe", help="Game to play")
 parser.add_argument("--allow-menu", action="store_true", help="Allow menu actions (Disabled by default)")
 parser.add_argument("--port", type=int, default=5555, help="Port for model server")
+parser.add_argument("--width", type=int, default=None, help="Override capture width")
+parser.add_argument("--height", type=int, default=None, help="Override capture height")
 
 args = parser.parse_args()
 
@@ -92,6 +94,8 @@ def preprocess_img(main_image):
 
 env = GamepadEnv(
     game=args.process,
+    image_width=args.width,
+    image_height=args.height,
     game_speed=1.0,
     env_fps=60,
     async_mode=True,