Preflight Planner Update
This commit is contained in:
293
src/utils/recorder.py
Normal file
293
src/utils/recorder.py
Normal file
@@ -0,0 +1,293 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Run Recorder — Captures logs, flight path, camera frames, and summary for each run.
|
||||
|
||||
Creates timestamped folders in /results with:
|
||||
flight_path.avi — Flight tracker video
|
||||
camera.avi — Downward camera video
|
||||
flight_path.png — Final flight tracker snapshot
|
||||
log.txt — Full console output
|
||||
summary.json — Run metadata and results
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import io
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
import cv2
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class RunRecorder:
|
||||
"""Records all outputs from a simulation run."""
|
||||
|
||||
def __init__(self, results_dir=None, fps=5):
|
||||
if results_dir is None:
|
||||
project_dir = Path(__file__).resolve().parent.parent
|
||||
results_dir = project_dir / "results"
|
||||
|
||||
results_dir = Path(results_dir)
|
||||
results_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Find next sequential run number
|
||||
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
|
||||
|
||||
self.fps = fps
|
||||
self.start_time = time.time()
|
||||
|
||||
# Log capture
|
||||
self._log_path = self.run_dir / "log.txt"
|
||||
self._log_file = open(self._log_path, "w")
|
||||
self._original_stdout = sys.stdout
|
||||
self._original_stderr = sys.stderr
|
||||
self._tee_stdout = _TeeWriter(sys.stdout, self._log_file)
|
||||
self._tee_stderr = _TeeWriter(sys.stderr, self._log_file)
|
||||
|
||||
# Video writers (initialized lazily on first frame)
|
||||
self._tracker_writer = None
|
||||
self._camera_writer = None
|
||||
self._tracker_size = None
|
||||
self._camera_size = None
|
||||
|
||||
# Frame counters
|
||||
self._tracker_frames = 0
|
||||
self._camera_frames = 0
|
||||
|
||||
# Snapshot storage
|
||||
self._last_tracker_frame = None
|
||||
self._last_camera_frame = None
|
||||
self._camera_snapshots = []
|
||||
|
||||
# Recording thread
|
||||
self._recording = False
|
||||
self._tracker_ref = None
|
||||
self._camera_ref = None
|
||||
self._record_thread = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Metadata
|
||||
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 start_logging(self):
|
||||
"""Redirect stdout/stderr to both console and log file."""
|
||||
sys.stdout = self._tee_stdout
|
||||
sys.stderr = self._tee_stderr
|
||||
|
||||
def stop_logging(self):
|
||||
"""Restore original stdout/stderr."""
|
||||
sys.stdout = self._original_stdout
|
||||
sys.stderr = self._original_stderr
|
||||
|
||||
def start_recording(self, tracker=None, camera=None):
|
||||
"""Start background recording of tracker and camera frames."""
|
||||
self._tracker_ref = tracker
|
||||
self._camera_ref = camera
|
||||
self._recording = True
|
||||
self._record_thread = threading.Thread(
|
||||
target=self._record_loop, daemon=True
|
||||
)
|
||||
self._record_thread.start()
|
||||
|
||||
def _record_loop(self):
|
||||
"""Periodically capture frames from tracker and camera."""
|
||||
interval = 1.0 / self.fps
|
||||
while self._recording:
|
||||
t0 = time.time()
|
||||
|
||||
# Capture tracker frame
|
||||
if self._tracker_ref is not None:
|
||||
try:
|
||||
frame = self._tracker_ref.draw()
|
||||
if frame is not None:
|
||||
self._write_tracker_frame(frame)
|
||||
self._last_tracker_frame = frame.copy()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Capture camera frame
|
||||
if self._camera_ref is not None:
|
||||
try:
|
||||
frame = self._camera_ref.frames.get("downward")
|
||||
if frame is not None:
|
||||
self._write_camera_frame(frame)
|
||||
self._last_camera_frame = frame.copy()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
elapsed = time.time() - t0
|
||||
sleep_time = max(0, interval - elapsed)
|
||||
time.sleep(sleep_time)
|
||||
|
||||
def _write_tracker_frame(self, frame):
|
||||
"""Write a frame to the tracker video."""
|
||||
h, w = frame.shape[:2]
|
||||
if self._tracker_writer is None:
|
||||
self._tracker_size = (w, h)
|
||||
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
||||
path = str(self.run_dir / "flight_path.avi")
|
||||
self._tracker_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
|
||||
self._tracker_writer.write(frame)
|
||||
self._tracker_frames += 1
|
||||
|
||||
def _write_camera_frame(self, frame):
|
||||
"""Write a frame to the camera video."""
|
||||
h, w = frame.shape[:2]
|
||||
if self._camera_writer is None:
|
||||
self._camera_size = (w, h)
|
||||
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
||||
path = str(self.run_dir / "camera.avi")
|
||||
self._camera_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
|
||||
self._camera_writer.write(frame)
|
||||
self._camera_frames += 1
|
||||
|
||||
def snapshot_camera(self, label="snapshot"):
|
||||
"""Save a named snapshot of the current camera frame."""
|
||||
if self._camera_ref is None:
|
||||
return
|
||||
frame = self._camera_ref.frames.get("downward")
|
||||
if frame is None:
|
||||
return
|
||||
idx = len(self._camera_snapshots)
|
||||
filename = f"camera_{idx:03d}_{label}.png"
|
||||
path = self.run_dir / filename
|
||||
cv2.imwrite(str(path), frame)
|
||||
self._camera_snapshots.append(filename)
|
||||
print(f"[REC] Snapshot: {filename}")
|
||||
|
||||
def save_summary(self, search_mode="", altitude=0, markers=None,
|
||||
landed=False, extra=None):
|
||||
"""Write the run summary JSON."""
|
||||
duration = time.time() - self.start_time
|
||||
mins = int(duration // 60)
|
||||
secs = int(duration % 60)
|
||||
|
||||
summary = {
|
||||
**self.metadata,
|
||||
"end_time": datetime.now().isoformat(),
|
||||
"duration_seconds": round(duration, 1),
|
||||
"duration_display": f"{mins}m {secs}s",
|
||||
"search_mode": search_mode,
|
||||
"altitude": altitude,
|
||||
"landed": landed,
|
||||
"markers_found": {},
|
||||
"recordings": {
|
||||
"log": "log.txt",
|
||||
"flight_path_video": "flight_path.avi",
|
||||
"flight_path_image": "flight_path.png",
|
||||
"camera_video": "camera.avi",
|
||||
"camera_snapshots": self._camera_snapshots,
|
||||
},
|
||||
"frame_counts": {
|
||||
"tracker": self._tracker_frames,
|
||||
"camera": self._camera_frames,
|
||||
},
|
||||
}
|
||||
|
||||
if markers:
|
||||
for mid, info in markers.items():
|
||||
pos = info.get('uav_position', {})
|
||||
summary["markers_found"][str(mid)] = {
|
||||
"x": round(pos.get('x', 0), 2),
|
||||
"y": round(pos.get('y', 0), 2),
|
||||
"distance": round(info.get('distance', 0), 2),
|
||||
}
|
||||
|
||||
if extra:
|
||||
summary.update(extra)
|
||||
|
||||
path = self.run_dir / "summary.json"
|
||||
with open(path, "w") as f:
|
||||
json.dump(summary, f, indent=2)
|
||||
print(f"[REC] Summary saved: {path}")
|
||||
|
||||
def stop(self):
|
||||
"""Stop recording and finalize all outputs."""
|
||||
self._recording = False
|
||||
if self._record_thread:
|
||||
self._record_thread.join(timeout=3.0)
|
||||
|
||||
# Save final flight path image
|
||||
if self._last_tracker_frame is not None:
|
||||
path = self.run_dir / "flight_path.png"
|
||||
cv2.imwrite(str(path), self._last_tracker_frame)
|
||||
print(f"[REC] Flight path saved: {path}")
|
||||
|
||||
# Save final camera frame
|
||||
if self._last_camera_frame is not None:
|
||||
path = self.run_dir / "camera_final.png"
|
||||
cv2.imwrite(str(path), self._last_camera_frame)
|
||||
|
||||
# Release video writers
|
||||
if self._tracker_writer:
|
||||
self._tracker_writer.release()
|
||||
if self._camera_writer:
|
||||
self._camera_writer.release()
|
||||
|
||||
# Stop log capture
|
||||
self.stop_logging()
|
||||
self._log_file.close()
|
||||
|
||||
duration = time.time() - self.start_time
|
||||
mins = int(duration // 60)
|
||||
secs = int(duration % 60)
|
||||
|
||||
# Print to original stdout since we stopped the tee
|
||||
self._original_stdout.write(
|
||||
f"\n[REC] Run recorded: {self.run_dir}\n"
|
||||
f"[REC] Duration: {mins}m {secs}s | "
|
||||
f"Tracker: {self._tracker_frames} frames | "
|
||||
f"Camera: {self._camera_frames} frames\n"
|
||||
)
|
||||
|
||||
|
||||
class _TeeWriter:
|
||||
"""Writes to both a stream and a file simultaneously."""
|
||||
|
||||
def __init__(self, stream, log_file):
|
||||
self._stream = stream
|
||||
self._log = log_file
|
||||
|
||||
def write(self, data):
|
||||
self._stream.write(data)
|
||||
try:
|
||||
# Strip ANSI escape codes for the log file
|
||||
clean = data
|
||||
self._log.write(clean)
|
||||
self._log.flush()
|
||||
except (ValueError, IOError):
|
||||
pass
|
||||
|
||||
def flush(self):
|
||||
self._stream.flush()
|
||||
try:
|
||||
self._log.flush()
|
||||
except (ValueError, IOError):
|
||||
pass
|
||||
|
||||
def fileno(self):
|
||||
return self._stream.fileno()
|
||||
|
||||
def isatty(self):
|
||||
return self._stream.isatty()
|
||||
Reference in New Issue
Block a user