Merge commit 'edebc795d9df8936dd9311b06bdb06de0fe5d9cb'

This commit is contained in:
Andrew Johnson
2025-11-09 13:31:30 +01:00
7 changed files with 180 additions and 33 deletions

1
.gitignore vendored
View File

@@ -87,3 +87,4 @@ cython_debug/
.devcontainer/devcontainer.json
output/
artifacts/

View File

@@ -58,7 +58,6 @@ If you prefer to set up the environment manually:
export PYTHONPATH=$PWD
```
## Running the simulation
To run the simulation, run the following command:

View File

@@ -20,12 +20,16 @@ tello.streamoff()
# Prepare directory to save
script_dir = os.path.dirname(__file__)
artifact_folder_path = os.path.join(script_dir, "../../artifacts/images")
artifact_folder_path = os.path.join(os.path.dirname(script_dir), "artifacts", "images")
os.makedirs(artifact_folder_path, exist_ok=True)
print("[Example] Saving captured picture to:", artifact_folder_path)
# Save the frame
save_path = os.path.join(artifact_folder_path, "picture.png")
cv2.imwrite(save_path, np.array(frame_read.frame))
# Land
tello.land()

View File

@@ -1,7 +1,7 @@
import os
import socket
from ursina import *
from cv2.typing import MatLike
import errno
from ursina import * # type: ignore
from time import time
import cv2
from ursina_adapter import UrsinaAdapter
@@ -17,10 +17,34 @@ class CommandServer:
self.stream_active = False
self.last_altitude = 0
self._recording_folder = "output/recordings"
self.server_socket = None
if not os.path.exists(self._recording_folder):
os.makedirs(self._recording_folder)
def check_port_available(self, port: int = 9999) -> bool:
"""
Check if the specified port is available.
Returns True if available, False if in use.
"""
try:
test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
test_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
test_socket.bind(('localhost', port))
test_socket.close()
return True
except OSError:
return False
def cleanup(self):
"""Clean up resources and close the server socket."""
if self.server_socket:
try:
self.server_socket.close()
print("[Command Listener] Server socket closed.")
except Exception as e:
print(f"[Command Listener] Error closing socket: {e}")
def streamon(self):
"""Start capturing screenshots and enable FPV video preview."""
if not self.stream_active:
@@ -60,16 +84,35 @@ class CommandServer:
"""
Listens for commands to send to the Simulator
"""
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.bind(('localhost', 9999)) # Port number for communication
server.listen(5)
print("[Command Listener] Listening on port 9999...")
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
while True:
conn, _ = server.accept()
data = conn.recv(1024).decode()
if data:
print(f"[Command Listener] Received command: {data}")
try:
self.server_socket.bind(('localhost', 9999)) # Port number for communication
self.server_socket.listen(5)
print("[Command Listener] Listening on port 9999...")
except OSError as e:
if e.errno == errno.EADDRINUSE: # Address already in use
print("\n" + "="*70)
print("ERROR: Port 9999 is already in use!")
print("="*70)
print("\nAnother instance of the simulator may be running.")
print("\nTo fix this, run one of these commands in your terminal:")
print(" macOS/Linux: lsof -ti:9999 | xargs kill -9")
print(" Windows: netstat -ano | findstr :9999")
print(" taskkill /PID <PID> /F")
print("\nOr simply restart your computer.")
print("="*70 + "\n")
raise
else:
raise
try:
while True:
conn, _ = self.server_socket.accept()
data = conn.recv(1024).decode()
if data:
print(f"[Command Listener] Received command: {data}")
if data == "connect":
self._ursina_adapter.connect()
@@ -202,13 +245,25 @@ class CommandServer:
conn.send(state.encode())
elif data == "get_latest_frame":
# Save the frame to disk first
frame_path = os.path.join(self._recording_folder, "latest_frame.png")
# Send frame data directly over TCP instead of using filesystem
if self._ursina_adapter.latest_frame is not None:
cv2.imwrite(frame_path, self._ursina_adapter.latest_frame)
conn.send(frame_path.encode())
# Encode frame as PNG in memory
success, buffer = cv2.imencode('.png', self._ursina_adapter.latest_frame)
if success:
# Send frame size first (4 bytes)
frame_data = buffer.tobytes()
frame_size = len(frame_data)
conn.send(frame_size.to_bytes(4, byteorder='big'))
# Then send the actual frame data
conn.send(frame_data)
print(f"[Frame Transfer] Sent {frame_size} bytes over TCP")
else:
# Send 0 size to indicate no frame
conn.send((0).to_bytes(4, byteorder='big'))
else:
conn.send(b"N/A")
# Send 0 size to indicate no frame available
conn.send((0).to_bytes(4, byteorder='big'))
elif data == "capture_frame":
self._ursina_adapter.capture_frame()
elif data.startswith("set_speed"):
@@ -221,4 +276,11 @@ class CommandServer:
elif data == "end":
self.end()
conn.close()
conn.close()
except KeyboardInterrupt:
print("\n[Command Listener] Shutting down...")
self.cleanup()
except Exception as e:
print(f"[Command Listener] Error: {e}")
self.cleanup()
raise

View File

@@ -1,21 +1,63 @@
from command_server import CommandServer
from ursina_adapter import UrsinaAdapter
import threading
import atexit
import signal
import sys
class TelloDroneSim:
def __init__(self):
self._ursina_adapter = UrsinaAdapter()
self._server = CommandServer(self._ursina_adapter)
# Register cleanup handlers
atexit.register(self.cleanup)
signal.signal(signal.SIGINT, self._signal_handler)
signal.signal(signal.SIGTERM, self._signal_handler)
def _signal_handler(self, signum, frame):
"""Handle termination signals gracefully."""
print("\n[Tello Sim] Received shutdown signal, cleaning up...")
self.cleanup()
sys.exit(0)
def cleanup(self):
"""Clean up resources."""
if hasattr(self, '_server'):
self._server.cleanup()
@property
def state(self):
return self._ursina_adapter
def start(self):
# Check if port is available before starting
if not self._server.check_port_available(9999):
print("\n" + "="*70)
print("ERROR: Cannot start simulator - Port 9999 is already in use!")
print("="*70)
print("\nAnother instance of the simulator may be running.")
print("\nTo fix this, run one of these commands in your terminal:")
print(" macOS/Linux: lsof -ti:9999 | xargs kill -9")
print(" Windows: netstat -ano | findstr :9999")
print(" taskkill /PID <PID> /F")
print("\nOr simply restart your computer.")
print("="*70 + "\n")
sys.exit(1)
server_thread = threading.Thread(target=self._server.listen)
server_thread.daemon = True
server_thread.start()
self._ursina_adapter.run()
try:
self._ursina_adapter.run()
except KeyboardInterrupt:
print("\n[Tello Sim] Interrupted, cleaning up...")
self.cleanup()
except Exception as e:
print(f"[Tello Sim] Error: {e}")
self.cleanup()
raise
def update(self) -> None:
self._ursina_adapter.tick()

View File

@@ -69,12 +69,53 @@ class TelloSimClient:
print(f"[Error] Unable to connect to the simulation at {self.host}:{self.port}")
def get_frame_read(self) -> BackgroundFrameRead:
frame_path = self._request_data('get_latest_frame')
if frame_path != "N/A" and os.path.exists(frame_path):
image = cv2.imread(frame_path)
if image is not None:
return BackgroundFrameRead(frame=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
"""Get the latest frame directly from the simulator over TCP."""
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.connect((self.host, self.port))
s.send(b'get_latest_frame')
# Receive frame size (4 bytes)
size_data = s.recv(4)
if len(size_data) != 4:
print("[Error] Failed to receive frame size")
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
frame_size = int.from_bytes(size_data, byteorder='big')
# If size is 0, no frame available
if frame_size == 0:
print("[Debug] No frame available from simulator")
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
# Receive the frame data
frame_data = b''
bytes_received = 0
while bytes_received < frame_size:
chunk = s.recv(min(4096, frame_size - bytes_received))
if not chunk:
break
frame_data += chunk
bytes_received += len(chunk)
# Decode the frame from PNG bytes
if len(frame_data) == frame_size:
nparr = np.frombuffer(frame_data, np.uint8)
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if image is not None:
# Return frame in BGR format (OpenCV's native format)
# Users should convert to RGB if needed for display
return BackgroundFrameRead(frame=image)
print("[Error] Failed to decode frame data")
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
except ConnectionRefusedError:
print(f"[Error] Unable to connect to the simulation at {self.host}:{self.port}")
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
except Exception as e:
print(f"[Error] Failed to get frame: {e}")
return BackgroundFrameRead(frame=np.zeros([360, 640, 3], dtype=np.uint8))
def _request_data(self, command):
try:

View File

@@ -828,7 +828,7 @@ class UrsinaAdapter():
def capture_frame(self):
"""Capture and save the latest FPV frame from update()"""
"""Capture the latest FPV frame. Optionally save to disk if save_frames_to_disk is True."""
if not self.stream_active:
print("[Capture] Stream not active. Cannot capture frame.")
return
@@ -837,11 +837,9 @@ class UrsinaAdapter():
print("[Capture] No latest frame available.")
return
frame_path = os.path.join(self.recording_folder, f"frame_{self.frame_count}.png")
cv2.imwrite(frame_path, self.latest_frame)
self.saved_frames.append(frame_path)
# Always increment frame count for tracking
self.frame_count += 1
print(f"[Capture] Screenshot {self.frame_count} saved: {frame_path}")
print(f"[Capture] Frame {self.frame_count} captured (memory only)")
def set_speed(self, x: int):
"""Set drone speed by adjusting acceleration force.
@@ -899,7 +897,7 @@ class UrsinaAdapter():
if self.bezier_mode:
t_now = time()
elapsed = t_now - self.bezier_start_time
elapsed = t_now - self.bezier_start_time # type: ignore
t = min(1.0, elapsed / self.bezier_duration)
# Bézier point