Add images and video conversion

samgeo/common.py

@@ -7,7 +7,7 @@
 import cv2
 import numpy as np
 from tqdm import tqdm
-
+from typing import List, Optional, Union
 import shapely
 import pyproj
 import rasterio
@@ -3184,3 +3184,142 @@ def choose_device(empty_cache: bool = True, quiet: bool = True) -> str:
                 "See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion."
             )
     return device
+
+
+def images_to_video(
+    images: Union[str, List[str]],
+    output_video: str,
+    fps: int = 30,
+    video_size: Optional[tuple] = None,
+) -> None:
+    """
+    Converts a series of images into a video. The input can be either a directory
+    containing the images or a list of image file paths.
+
+    Args:
+        images (Union[str, List[str]]): A directory containing images or a list
+            of image file paths.
+        output_video (str): The filename of the output video (e.g., 'output.mp4').
+        fps (int, optional): Frames per second for the output video. Default is 30.
+        video_size (Optional[tuple], optional): The size (width, height) of the
+            video. If not provided, the size of the first image is used.
+
+    Raises:
+        ValueError: If the provided path is not a directory, if the images list
+            is empty, or if the first image cannot be read.
+
+    Example usage:
+        images_to_video('path_to_image_directory', 'output_video.mp4', fps=30, video_size=(1280, 720))
+        images_to_video(['image1.jpg', 'image2.jpg', 'image3.jpg'], 'output_video.mp4', fps=30)
+    """
+    if isinstance(images, str):
+        if not os.path.isdir(images):
+            raise ValueError(f"The provided path {images} is not a valid directory.")
+
+        # Get all image files in the directory (sorted by filename)
+        images = [
+            os.path.join(images, img)
+            for img in sorted(os.listdir(images))
+            if img.endswith((".jpg", ".png"))
+        ]
+
+    if not isinstance(images, list) or not images:
+        raise ValueError(
+            "The images parameter should either be a non-empty list of image paths or a valid directory."
+        )
+
+    # Read the first image to get the dimensions if video_size is not provided
+    first_image_path = images[0]
+    frame = cv2.imread(first_image_path)
+
+    if frame is None:
+        raise ValueError(f"Error reading the first image {first_image_path}")
+
+    if video_size is None:
+        height, width, _ = frame.shape
+        video_size = (width, height)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Define the codec for mp4
+    video_writer = cv2.VideoWriter(output_video, fourcc, fps, video_size)
+
+    for image_path in images:
+        frame = cv2.imread(image_path)
+        if frame is None:
+            print(f"Warning: Could not read image {image_path}. Skipping.")
+            continue
+
+        if video_size != (frame.shape[1], frame.shape[0]):
+            frame = cv2.resize(frame, video_size)
+
+        video_writer.write(frame)
+
+    video_writer.release()
+    print(f"Video saved as {output_video}")
+
+
+def video_to_images(
+    video_path: str,
+    output_dir: str,
+    frame_rate: Optional[int] = None,
+    prefix: str = "frame_",
+) -> None:
+    """
+    Converts a video into a series of images. Each frame of the video is saved as an image.
+
+    Args:
+        video_path (str): The path to the video file.
+        output_dir (str): The directory where the images will be saved.
+        frame_rate (Optional[int], optional): The number of frames to save per second of video.
+            If None, all frames will be saved. Defaults to None.
+        prefix (str, optional): The prefix for the output image filenames. Defaults to 'frame_'.
+
+    Raises:
+        ValueError: If the video file cannot be read or if the output directory is invalid.
+
+    Example usage:
+        video_to_images('input_video.mp4', 'output_images', frame_rate=1, prefix='image_')
+    """
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Error opening video file {video_path}")
+
+    # Get video properties
+    video_fps = int(cap.get(cv2.CAP_PROP_FPS))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_rate = (
+        frame_rate if frame_rate else video_fps
+    )  # Default to original FPS if not provided
+
+    # Calculate the number of digits based on the total frames (e.g., if total frames are 1000, width = 4)
+    num_digits = len(str(total_frames))
+
+    print(f"Video FPS: {video_fps}")
+    print(f"Total Frames: {total_frames}")
+    print(f"Saving every {video_fps // frame_rate} frame(s)")
+
+    frame_count = 0
+    saved_frame_count = 0
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Save frames based on frame_rate
+        if frame_count % (video_fps // frame_rate) == 0:
+            img_path = os.path.join(
+                output_dir, f"{prefix}{saved_frame_count:0{num_digits}d}.jpg"
+            )
+            cv2.imwrite(img_path, frame)
+            saved_frame_count += 1
+            # print(f"Saved {img_path}")
+
+        frame_count += 1
+
+    # Release the video capture object
+    cap.release()
+    print(f"Finished saving {saved_frame_count} images to {output_dir}")