Updated frame saving and camera mode setting.

Author: brittleru

drone_base/config/video.py

@@ -1,7 +1,7 @@
 class VideoConfig:
     """Configuration class for the drone video stream with validation"""
 
-    def __init__(self, width: int = 640, height: int = 360, cam_mode: str | None = None):
+    def __init__(self, width: int = 640, height: int = 360, cam_mode: str | None = None, save_extension: str = "png"):
         """
         Sets the width and height of the video stream. Computes the center of the frame.
 
@@ -12,12 +12,15 @@ def __init__(self, width: int = 640, height: int = 360, cam_mode: str | None = N
         """
         if width <= 0 or height <= 0:
             raise ValueError("Video width and height must be greater than zero.")
+        if save_extension not in ("png", "jpg", "jpeg"):
+            raise ValueError(f"Save extension must be 'png' or 'jpg' or 'jpeg'. Provided {save_extension}")
 
         self.width: int = width
         self.height: int = height
         self.frame_center_x, self.frame_center_y = self.__compute_frame_center()
         self.frame_area = self.__compute_frame_area()
         self.max_queue_size = 30
+        self.save_extension = save_extension
 
         if cam_mode not in ("photo", "recording", None):
             raise ValueError(

drone_base/stream/base_streaming_controller.py

@@ -22,20 +22,22 @@
 class BaseStreamingController(ABC):
     """Base class for all streaming controllers with common functionality."""
 
-    def __init__(
+    def __init__( # noqa: PLR0913
             self,
             ip: DroneIp,
             processor_class: type[BaseVideoProcessor],
             speed: int = 35,
             log_path: str | Path | None = None,
             results_path: str | Path | None = None,
+            video_config: VideoConfig | None = None,
     ):
         self.start_time = date_time_now_to_file_name()
         if log_path is not None:
             log_path = Path(log_path) / self.start_time
         if results_path is not None:
             results_path = Path(results_path) / self.start_time
-        self.video_config = VideoConfig(width=640, height=360, cam_mode="photo")
+
+        self.video_config = video_config if video_config is not None else VideoConfig()
 
         self.drone = olympe.Drone(ip)
         self.drone_commander = DroneCommander(self.drone, logger_dir=log_path)

drone_base/stream/base_video_processor.py

@@ -34,7 +34,9 @@ def __init__(
         self.is_frame_saved = frame_save_dir is not None
         if self.is_frame_saved:
             save_path = Path(frame_save_dir) / "frames"
-            self.frame_saver = BufferedFrameSaver(output_dir=save_path, logger_dir=logger_dir)
+            self.frame_saver = BufferedFrameSaver(
+                output_dir=save_path, logger_dir=logger_dir, save_extension=video_config.save_extension
+            )
 
         if logger_dir is not None:
             logger_dir = Path(logger_dir) / f"{self.__class__.__name__}.log"

drone_base/stream/saving/frame_saver.py

@@ -7,8 +7,9 @@
 
 
 class BufferedFrameSaver:
-    def __init__(self, output_dir: str | Path, logger_dir: str | Path | None = None):
+    def __init__(self, output_dir: str | Path, logger_dir: str | Path | None = None, save_extension: str = "png"):
         self.output_dir = Path(output_dir)
+        self.save_extension = save_extension
         self.output_dir.mkdir(parents=True, exist_ok=True)
         self.frames: list[tuple[np.ndarray, float]] = []
 
@@ -25,7 +26,7 @@ def save_all(self):
         self.logger.info("Saving all frames to disk")
         for index, (frame, timestamp) in enumerate(self.frames):
             timestamp_ms = int(timestamp * 1000)
-            output_path = self.output_dir / f"frame_{index:06d}_{timestamp_ms}.png"
+            output_path = self.output_dir / f"frame_{index:06d}_{timestamp_ms}.{self.save_extension}"
             cv2.imwrite(str(output_path), frame)
             if index % 100 == 0:
                 self.logger.info("Saved %s/%s frames", f"{index:06d}", len(self.frames))