georef_capture.py

import cv2
import threading
from pymavlink import mavutil
import argparse
import os
import datetime
import yaml
import copy
import dataclasses
from frame_selector import FrameSelector
from models import DroneData
import logging


STREAM_PIPELINE = " out. ! queue ! rtph264pay config-interval=1 pt=96 ! udpsink host={address} port={port} sync=false async=false"


class Capturer:
    def __init__(self, args):
        with open(args.config, "r") as f:
            self.config = yaml.safe_load(f)
        logging.info("Loaded config:")
        logging.info("\n" + yaml.dump(self.config, sort_keys=False))
        self.preview = args.preview
        self.no_tele = args.no_tele
        self.stream_ip = args.stream_ip
        self.last_drone_data = None
        self.telems = []
        self.cap = None
        self.writer = None
        self.mav_listener = None
        self.running = False

        self.output_dir = datetime.datetime.now().strftime("data/%Y-%m-%d_%H-%M-%S")
        os.makedirs(self.output_dir, exist_ok=True)
        base = os.path.basename("out.mp4")
        self.video_filename = os.path.join(self.output_dir, base)
        frames_dir = os.path.join(self.output_dir, "frames")
        os.makedirs(frames_dir, exist_ok=True)
        self.frame_selector = FrameSelector(
            frames_dir,
            self.config.get("gcs_ip"),
            self.config.get("select_nth_frame"),
            self.config.get("select_on_request", False)
        )

    def run(self):
        self.running = True
        if not self.no_tele:
            self.mav_listener = threading.Thread(target=self.mavlink_listener)
            self.mav_listener.start()
        else:
            logging.info(
                "Telemetry disabled, not starting MAVLink listener, georef frames will not be saved."
            )
        self.video_capture()
        self.finish()

    def finish(self):
        self.running = False
        self.frame_selector.finish_saving()
        if not self.no_tele:
            logging.info(f"Captured {len(self.telems)} telemetry points.")
            yaml_path = os.path.join(self.output_dir, "telemetry.yaml")
            with open(yaml_path, "w") as f:
                yaml.dump([dataclasses.asdict(t) for t in self.telems], f)
            logging.info(f"Telemetry saved to {yaml_path}")

    def mavlink_listener(self):
        logging.info("Mavlink listener started")
        master = mavutil.mavlink_connection(
            self.config.get("connection_string"), baud=self.config.get("baud_rate")
        )

        logging.info("Waiting for heartbeat...")
        master.wait_heartbeat()
        logging.info("Heartbeat received!")
        while self.running:
            msg = master.recv_match(
                type=["GLOBAL_POSITION_INT", "ATTITUDE", "CAMERA_IMAGE_CAPTURED", "CAMERA_TRIGGER"], blocking=True
            )
            if not msg:
                continue
            if self.last_drone_data is None:
                self.last_drone_data = DroneData()

            if msg.get_type() == "GLOBAL_POSITION_INT":
                self.last_drone_data.lat = msg.lat / 1e7
                self.last_drone_data.lon = msg.lon / 1e7
                self.last_drone_data.alt = msg.alt / 1000.0  # in meters
                self.last_drone_data.rel_alt = msg.relative_alt / 1000.0

            elif msg.get_type() == "ATTITUDE":
                self.last_drone_data.roll = round(msg.roll, 5)
                self.last_drone_data.pitch = round(msg.pitch, 5)
                self.last_drone_data.yaw = round(msg.yaw, 5)
            elif msg.get_type() == "CAMERA_IMAGE_CAPTURED" or msg.get_type() == "CAMERA_TRIGGER":
                logging.debug(f"Camera msg received {msg.get_type()}.")
                self.frame_selector.request_saving()

    def prepare_gst_writer(self):
        ret, frame = self.cap.read()
        if not ret:
            logging.error("Error: Could not read frame for video writer initialization")
            self.cap.release()
            return

        height, width = frame.shape[:2]
        gst_writer_pipeline = self.config.get("gst_writer_pipeline")
        gst_writer_pipeline = gst_writer_pipeline.format(
            output_file=self.video_filename
        )
        if self.stream_ip:
            address, port = self.stream_ip.split(":")
            gst_writer_pipeline += STREAM_PIPELINE.format(address=address, port=port)
        logging.info(f"GStreamer writer pipeline: {gst_writer_pipeline}")
        self.writer = cv2.VideoWriter(
            gst_writer_pipeline,
            cv2.CAP_GSTREAMER,
            0,
            self.config.get("fps"),
            (width, height),
        )
        if self.writer is None or not self.writer.isOpened():
            logging.error(
                f"Error: Could not open video file {self.video_filename} for writing (H264). Try .mkv"
            )
            self.writer = None

    def video_capture(self):
        self.cap = cv2.VideoCapture(
            self.config.get("gst_capture_pipeline"),
            cv2.CAP_GSTREAMER,
        )
        self.prepare_gst_writer()
        frames_num = 0
        if not self.cap.isOpened():
            logging.error("Error: Unable to open camera")
            exit()
        try:
            while True:
                ret, frame = self.cap.read()
                if ret:
                    if self.preview:
                        cv2.imshow("Frame", frame)
                    if self.writer:
                        if (self.last_drone_data and self.last_drone_data.is_initialized()) or self.no_tele:
                            self.telems.append(copy.deepcopy(self.last_drone_data))
                            self.writer.write(frame)
                            self.frame_selector.take_frame(frame, self.last_drone_data)
                            frames_num += 1
                        else:
                            logging.warning("Warning: No telem, skipping frame.")
                else:
                    logging.error("Error: Could not read frame")
                    break

                if self.preview and cv2.waitKey(1) == ord("q"):
                    break
        except KeyboardInterrupt:
            logging.info("\nKeyboardInterrupt received in video loop. Exiting...")
        finally:
            self.cap.release()
            if self.writer:
                self.writer.release()
            logging.info(f"Total frames written: {frames_num}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Georeferenced video capture")
    parser.add_argument(
        "-p",
        "--preview",
        action="store_true",
        help="Show video preview window (default: off)",
    )
    parser.add_argument(
        "--no-tele",
        action="store_true",
        help="Disable telemetry data saving (default: off)",
    )
    parser.add_argument(
        "--stream-ip",
        metavar="STREAM_IP",
        help="UDP address for video streaming",
    )
    parser.add_argument(
        "-c",
        "--config",
        metavar="CONFIG_FILE",
        type=str,
        required=True,
        help="YAML config file path to use (required)",
    )
    parser.add_argument(
        "-l",
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Set the logging level (default: INFO)",
    )
    args = parser.parse_args()
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(levelname)s: %(message)s",
    )
    capturer = Capturer(args)
    capturer.run()