Cloud Saves Update
This commit is contained in:
@@ -1,13 +1,7 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run recorder — captures simulation video and logs to a results folder.
|
||||
"""Run recorder — captures simulation video and logs, and uploads to local REST API backend.
|
||||
|
||||
Each run gets a folder like results/run_1/ containing:
|
||||
- log.txt Process logs (stdout/stderr)
|
||||
- flight_path.avi Flight tracker video
|
||||
- flight_path.png Final flight path image
|
||||
- camera.avi Downward camera video
|
||||
- camera_*.png Camera snapshots at key moments
|
||||
- gazebo.avi Gazebo simulation window recording (via xwd)
|
||||
Each run gets a temporary folder, and all resources are uploaded directly to the backend.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -19,8 +13,11 @@ import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import re
|
||||
import tempfile
|
||||
import requests
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
import yaml
|
||||
|
||||
try:
|
||||
import PIL.Image
|
||||
@@ -28,32 +25,57 @@ try:
|
||||
except ImportError:
|
||||
HAS_PIL = False
|
||||
|
||||
API_URL = os.environ.get("API_URL", "https://simlink.sirblob.co")
|
||||
|
||||
class RunRecorder:
|
||||
def __init__(self, results_dir=None, fps=5):
|
||||
if results_dir is None:
|
||||
project_dir = Path(__file__).resolve().parent.parent.parent
|
||||
results_dir = project_dir / "results"
|
||||
|
||||
results_dir = Path(results_dir)
|
||||
results_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
existing = [d.name for d in results_dir.iterdir()
|
||||
if d.is_dir() and d.name.startswith("run_")]
|
||||
nums = []
|
||||
for name in existing:
|
||||
try:
|
||||
nums.append(int(name.split("_")[1]))
|
||||
except (IndexError, ValueError):
|
||||
pass
|
||||
run_num = max(nums, default=0) + 1
|
||||
|
||||
self.run_dir = results_dir / f"run_{run_num}"
|
||||
self.run_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.run_num = run_num
|
||||
|
||||
project_dir = Path(__file__).resolve().parent.parent.parent
|
||||
self.fps = fps
|
||||
self.start_time = time.time()
|
||||
|
||||
self.search_start = 0.0
|
||||
self.search_duration = 0.0
|
||||
|
||||
# Gather config data as JSON
|
||||
config_data = {}
|
||||
for cfg in ["search.yaml", "ugv.yaml", "uav.yaml"]:
|
||||
p = project_dir / "config" / cfg
|
||||
if p.exists():
|
||||
try:
|
||||
with open(p, 'r') as f:
|
||||
data = yaml.safe_load(f)
|
||||
if data:
|
||||
key = cfg.replace('.yaml', '')
|
||||
config_data[key] = data
|
||||
except Exception as e:
|
||||
print(f"[REC] Failed to load config {cfg}: {e}")
|
||||
|
||||
if "ugv" in config_data and "topics" in config_data["ugv"]:
|
||||
del config_data["ugv"]["topics"]
|
||||
|
||||
if "uav" in config_data and "connection" in config_data["uav"]:
|
||||
del config_data["uav"]["connection"]
|
||||
|
||||
self.sim_id = 0
|
||||
self.sim_name = "simulation_unknown"
|
||||
try:
|
||||
resp = requests.post(f"{API_URL}/api/simulations/create", json=config_data, timeout=10)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
self.sim_id = data.get("id", 0)
|
||||
self.sim_name = data.get("name", f"simulation_{self.sim_id}")
|
||||
except Exception as e:
|
||||
print(f"[REC] API create failed: {e}")
|
||||
|
||||
if results_dir:
|
||||
self.run_dir = Path(results_dir) / self.sim_name
|
||||
else:
|
||||
self.run_dir = project_dir / "results" / self.sim_name
|
||||
self.run_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.run_num = self.sim_id
|
||||
|
||||
print(f"[REC] Recording local results to {self.run_dir.name} and API ({self.sim_name}, ID: {self.sim_id})")
|
||||
self._log_path = self.run_dir / "log.txt"
|
||||
self._log_file = open(self._log_path, "w")
|
||||
self._original_stdout = sys.stdout
|
||||
@@ -87,18 +109,23 @@ class RunRecorder:
|
||||
self._gazebo_wid = None
|
||||
self._find_gazebo_window()
|
||||
|
||||
self.metadata = {
|
||||
"run": run_num,
|
||||
"start_time": datetime.now().isoformat(),
|
||||
"run_dir": str(self.run_dir),
|
||||
}
|
||||
|
||||
print(f"[REC] Recording to: {self.run_dir}")
|
||||
def _upload_file(self, path, filename):
|
||||
if not self.sim_id:
|
||||
return
|
||||
try:
|
||||
with open(path, 'rb') as f:
|
||||
resp = requests.post(
|
||||
f"{API_URL}/api/simulations/{self.sim_id}/upload",
|
||||
files={"file": (filename, f)},
|
||||
timeout=60
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"[REC] Upload failed for {filename}: {e}")
|
||||
|
||||
def _find_gazebo_window(self):
|
||||
"""Attempts to find the Gazebo window ID using xwininfo."""
|
||||
try:
|
||||
# Find window ID
|
||||
cmd = "xwininfo -root -tree | grep -i 'gazebo\|gz sim'"
|
||||
output = subprocess.check_output(cmd, shell=True).decode('utf-8')
|
||||
lines = output.strip().split('\n')
|
||||
@@ -106,24 +133,24 @@ class RunRecorder:
|
||||
print("[REC] Gazebo window not found")
|
||||
return
|
||||
|
||||
# Pick the first one
|
||||
wid_line = lines[0]
|
||||
wid_match = re.search(r'(0x[0-9a-fA-F]+)', wid_line)
|
||||
if not wid_match:
|
||||
return
|
||||
self._gazebo_wid = wid_match.group(1)
|
||||
print(f"[REC] Found Gazebo window ID: {self._gazebo_wid}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[REC] Error finding Gazebo window: {e}")
|
||||
|
||||
def set_phase(self, phase: str):
|
||||
# No-op since CSV logging is removed
|
||||
pass
|
||||
if phase == "takeoff":
|
||||
self.start_time = time.time()
|
||||
elif phase == "search":
|
||||
self.search_start = time.time()
|
||||
elif phase != "search" and phase != "takeoff" and self.search_start > 0 and self.search_duration == 0:
|
||||
self.search_duration = time.time() - self.search_start
|
||||
|
||||
def log_position(self, uav_x=0, uav_y=0, uav_alt=0, uav_heading=0,
|
||||
ugv_x=0, ugv_y=0):
|
||||
# No-op since CSV logging is removed
|
||||
def log_position(self, uav_x=0, uav_y=0, uav_alt=0, uav_heading=0, ugv_x=0, ugv_y=0):
|
||||
pass
|
||||
|
||||
def start_logging(self):
|
||||
@@ -139,13 +166,10 @@ class RunRecorder:
|
||||
self._camera_ref = camera
|
||||
self._recording = True
|
||||
|
||||
# Re-check Gazebo window
|
||||
if not self._gazebo_wid:
|
||||
self._find_gazebo_window()
|
||||
|
||||
self._record_thread = threading.Thread(
|
||||
target=self._record_loop, daemon=True
|
||||
)
|
||||
self._record_thread = threading.Thread(target=self._record_loop, daemon=True)
|
||||
self._record_thread.start()
|
||||
|
||||
def _record_loop(self):
|
||||
@@ -171,31 +195,24 @@ class RunRecorder:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Gazebo Capture via xwd
|
||||
if self._gazebo_wid and HAS_PIL:
|
||||
try:
|
||||
# Capture window to stdout using xwd
|
||||
proc = subprocess.run(
|
||||
['xwd', '-id', self._gazebo_wid, '-out', '/dev/stdout'],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE
|
||||
)
|
||||
if proc.returncode == 0:
|
||||
# Parse XWD using PIL
|
||||
img_pil = PIL.Image.open(io.BytesIO(proc.stdout))
|
||||
# Convert to OpenCV (RGB -> BGR)
|
||||
img_np = np.array(img_pil)
|
||||
# PIL opens XWD as RGB usually, check mode
|
||||
if img_pil.mode == 'RGB':
|
||||
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
elif img_pil.mode == 'RGBA':
|
||||
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGR)
|
||||
else:
|
||||
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) # Fallback
|
||||
|
||||
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
self._write_gazebo_frame(img_bgr)
|
||||
except Exception:
|
||||
# xwd failed or window closed
|
||||
pass
|
||||
|
||||
elapsed = time.time() - t0
|
||||
@@ -224,7 +241,6 @@ class RunRecorder:
|
||||
|
||||
def _write_gazebo_frame(self, frame):
|
||||
h, w = frame.shape[:2]
|
||||
# Ensure dimensions are even
|
||||
if w % 2 != 0: w -= 1
|
||||
if h % 2 != 0: h -= 1
|
||||
frame = frame[:h, :w]
|
||||
@@ -235,7 +251,6 @@ class RunRecorder:
|
||||
path = str(self.run_dir / "gazebo.avi")
|
||||
self._gazebo_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
|
||||
elif (w, h) != self._gazebo_size:
|
||||
# Handle resize? Just skip for now or resize
|
||||
frame = cv2.resize(frame, self._gazebo_size)
|
||||
|
||||
self._gazebo_writer.write(frame)
|
||||
@@ -253,20 +268,17 @@ class RunRecorder:
|
||||
cv2.imwrite(str(path), frame)
|
||||
self._camera_snapshots.append(filename)
|
||||
print(f"[REC] Snapshot: {filename}")
|
||||
threading.Thread(target=self._upload_file, args=(path, filename), daemon=True).start()
|
||||
|
||||
def snapshot_tracker(self, label="tracker"):
|
||||
"""Save a snapshot of the current flight tracker."""
|
||||
if self._last_tracker_frame is not None:
|
||||
filename = f"tracker_{label}.png"
|
||||
path = self.run_dir / filename
|
||||
cv2.imwrite(str(path), self._last_tracker_frame)
|
||||
print(f"[REC] Tracker snapshot: {filename}")
|
||||
threading.Thread(target=self._upload_file, args=(path, filename), daemon=True).start()
|
||||
|
||||
def save_summary(self, search_mode="", altitude=0, markers=None,
|
||||
landed=False, ugv_dispatched=False,
|
||||
ugv_target=None, extra=None):
|
||||
# Summary JSON file saving moved/removed as per request
|
||||
# Still useful to see basic stats in terminal?
|
||||
def save_summary(self, search_mode="", altitude=0, markers=None, landed=False, ugv_dispatched=False, ugv_target=None, extra=None):
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
@@ -275,32 +287,49 @@ class RunRecorder:
|
||||
self._record_thread.join(timeout=3.0)
|
||||
|
||||
if self._last_tracker_frame is not None:
|
||||
path = self.run_dir / "flight_path.png"
|
||||
filename = "flight_path.png"
|
||||
path = self.run_dir / filename
|
||||
cv2.imwrite(str(path), self._last_tracker_frame)
|
||||
print(f"[REC] Flight path saved: {path}")
|
||||
self._upload_file(path, filename)
|
||||
|
||||
if self._last_camera_frame is not None:
|
||||
path = self.run_dir / "camera_final.png"
|
||||
filename = "camera_final.png"
|
||||
path = self.run_dir / filename
|
||||
cv2.imwrite(str(path), self._last_camera_frame)
|
||||
self._upload_file(path, filename)
|
||||
|
||||
if self._tracker_writer:
|
||||
self._tracker_writer.release()
|
||||
self._upload_file(self.run_dir / "flight_path.avi", "flight_path.avi")
|
||||
if self._camera_writer:
|
||||
self._camera_writer.release()
|
||||
self._upload_file(self.run_dir / "camera.avi", "camera.avi")
|
||||
if self._gazebo_writer:
|
||||
self._gazebo_writer.release()
|
||||
self._upload_file(self.run_dir / "gazebo.avi", "gazebo.avi")
|
||||
|
||||
self.stop_logging()
|
||||
self._log_file.close()
|
||||
self._upload_file(self._log_path, "log.txt")
|
||||
|
||||
duration = time.time() - self.start_time
|
||||
mins = int(duration // 60)
|
||||
secs = int(duration % 60)
|
||||
|
||||
if self.sim_id:
|
||||
try:
|
||||
requests.put(
|
||||
f"{API_URL}/api/simulations/{self.sim_id}/time",
|
||||
json={"search_time": self.search_duration, "total_time": duration},
|
||||
timeout=5
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[REC] PUT time failed: {e}")
|
||||
|
||||
self._original_stdout.write(
|
||||
f"\n[REC] ═══════════════════════════════════════\n"
|
||||
f"[REC] Run #{self.run_num} recorded\n"
|
||||
f"[REC] Dir: {self.run_dir}\n"
|
||||
f"[REC] Cloud Upload Complete (ID: {self.sim_id})\n"
|
||||
f"[REC] Duration: {mins}m {secs}s\n"
|
||||
f"[REC] Tracker: {self._tracker_frames} frames\n"
|
||||
f"[REC] Camera: {self._camera_frames} frames\n"
|
||||
@@ -308,7 +337,6 @@ class RunRecorder:
|
||||
f"[REC] ═══════════════════════════════════════\n"
|
||||
)
|
||||
|
||||
|
||||
class _TeeWriter:
|
||||
def __init__(self, stream, log_file):
|
||||
self._stream = stream
|
||||
|
||||
Reference in New Issue
Block a user