diff --git a/tello_sim/command_server.py b/tello_sim/command_server.py index d3afb42..e791fc7 100644 --- a/tello_sim/command_server.py +++ b/tello_sim/command_server.py @@ -18,7 +18,6 @@ class CommandServer: def __init__(self, ursina_adapter: UrsinaAdapter): self._ursina_adapter = ursina_adapter - self.latest_frame = None self.stream_active = False self.server_socket = None diff --git a/tello_sim/ursina_adapter.py b/tello_sim/ursina_adapter.py index 71c5836..dc98687 100644 --- a/tello_sim/ursina_adapter.py +++ b/tello_sim/ursina_adapter.py @@ -34,7 +34,6 @@ from time import sleep, time import time as time_module # ursina publishes the per-frame delta as time.dt on the module import traceback -from cv2.typing import MatLike logger = logging.getLogger(__name__) @@ -1419,13 +1418,6 @@ def emergency(self) -> None: else: print("Drone is already on the ground") - def get_latest_frame(self) -> MatLike: - """Return the latest frame directly""" - if self.latest_frame is None: - raise Exception("No latest frame available.") - return cv2.cvtColor(self.latest_frame, cv2.COLOR_BGR2RGB) - - def capture_frame(self): """Capture the latest FPV frame. Optionally save to disk if save_frames_to_disk is True.""" if not self.stream_active: @@ -1515,8 +1507,15 @@ def _tick_impl(self) -> None: # streamed frame. cvtColor and cv2.flip each return one fresh # array, so this path makes two allocations instead of four and # drops the (undeclared, transitive) Pillow dependency. + # + # Convert straight to RGB (not BGR): get_frame_read must hand + # back RGB to match djitellopy, and the frame only ever leaves + # here through a cv2.imencode('.png') -> cv2.imdecode round + # trip, which is an identity on the channel order. Storing RGB + # lets the client return the decoded frame as-is instead of + # running a second full-frame COLOR_BGR2RGB on every read. arr = np.frombuffer(pixel_data, np.uint8).reshape(height, width, 4) # type: ignore - frame = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR) + frame = cv2.cvtColor(arr, cv2.COLOR_RGBA2RGB) frame = cv2.flip(frame, 0) # cv2.flip returns a freshly allocated array that nothing else diff --git a/tello_sim_client.py b/tello_sim_client.py index 3a78d49..ae0c3b5 100644 --- a/tello_sim_client.py +++ b/tello_sim_client.py @@ -45,11 +45,13 @@ def get_frame_read(self) -> BackgroundFrameRead: print("[Error] Failed to decode frame data") return blank - # imdecode yields BGR; convert to RGB so the simulator's frames match - # the real Tello's djitellopy interface, which returns RGB. Consumers - # using OpenCV (cv2.imshow / imwrite / cvtColor to HSV) convert - # RGB->BGR themselves, exactly as they must for a real drone. - return BackgroundFrameRead(frame=cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) + # The simulator already captures the framebuffer as RGB, and the + # imencode('.png') -> imdecode round trip preserves channel order, so the + # decoded frame is RGB — matching the real Tello's djitellopy interface, + # which returns RGB. No client-side color conversion is needed. Consumers + # using OpenCV (cv2.imshow / imwrite / cvtColor to HSV) convert RGB->BGR + # themselves, exactly as they must for a real drone. + return BackgroundFrameRead(frame=image) def get_battery(self): return self._conn.request('get_battery') diff --git a/tests/test_fpv_frame_color.py b/tests/test_fpv_frame_color.py new file mode 100644 index 0000000..b5a6b2c --- /dev/null +++ b/tests/test_fpv_frame_color.py @@ -0,0 +1,70 @@ +"""Locks in the color-channel invariant the FPV frame path depends on. + +The simulator captures the framebuffer as RGB (cv2.COLOR_RGBA2RGB), the +command server ships it as a PNG via cv2.imencode, and the client rebuilds it +with cv2.imdecode and hands it back *without* any further color conversion +(tello_sim_client.TelloSimClient.get_frame_read). + +That is only correct because the encode and decode are BOTH OpenCV: imencode +treats its input array as BGR and writes a standard RGB PNG (swapping channels +internally), and imdecode reads that RGB PNG straight back into a BGR array +(swapping again). The two swaps cancel, so the imencode -> imdecode round trip +is an *identity on the array bytes* no matter what the channels actually mean. +It is the pairing that makes it channel-agnostic, not imdecode on its own +(which is indeed BGR-by-convention in isolation). + +This module proves that empirically, with no test framework required: + + python tests/test_fpv_frame_color.py + +It is also plain-`pytest` discoverable if the repo adopts one later. +""" +import numpy as np +import cv2 + + +def _sim_capture_rgb(rgba: np.ndarray) -> np.ndarray: + """Simulator's capture conversion (ursina_adapter._tick_impl).""" + return cv2.cvtColor(rgba, cv2.COLOR_RGBA2RGB) + + +def _png_roundtrip(frame: np.ndarray) -> np.ndarray: + """Server encode (command_server get_latest_frame) + client decode.""" + ok, buffer = cv2.imencode(".png", frame) + assert ok, "cv2.imencode failed" + return cv2.imdecode(np.frombuffer(buffer.tobytes(), np.uint8), cv2.IMREAD_COLOR) + + +def test_client_receives_rgb_without_conversion() -> None: + """The frame the client decodes already equals the RGB the sim captured, + so get_frame_read needs no client-side cvtColor.""" + rng = np.random.default_rng(0) + rgba = rng.integers(0, 256, size=(360, 640, 4), dtype=np.uint8) + rgba[:, :, 3] = 255 # opaque; the alpha channel is dropped on capture + + captured_rgb = _sim_capture_rgb(rgba) + client_frame = _png_roundtrip(captured_rgb) + + assert np.array_equal(client_frame, captured_rgb) + + +def test_new_pipeline_is_byte_identical_to_old() -> None: + """The RGB output handed to callers is byte-identical to the previous + BGR-capture + client BGR2RGB pipeline — same frames, one fewer conversion.""" + rng = np.random.default_rng(1) + rgba = rng.integers(0, 256, size=(360, 640, 4), dtype=np.uint8) + rgba[:, :, 3] = 255 + + # Old: capture BGR, PNG round trip, then client converts BGR -> RGB. + old = cv2.cvtColor(_png_roundtrip(cv2.cvtColor(rgba, cv2.COLOR_RGBA2BGR)), + cv2.COLOR_BGR2RGB) + # New: capture RGB, PNG round trip, client returns it as-is. + new = _png_roundtrip(cv2.cvtColor(rgba, cv2.COLOR_RGBA2RGB)) + + assert np.array_equal(old, new) + + +if __name__ == "__main__": + test_client_receives_rgb_without_conversion() + test_new_pipeline_is_byte_identical_to_old() + print("OK: FPV frame color round-trip invariants hold")