sam2-repo/demo/backend/server/data/loader.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import shutil
import subprocess
from glob import glob
from pathlib import Path
from typing import Dict, Optional

import imagesize
from app_conf import GALLERY_PATH, POSTERS_PATH, POSTERS_PREFIX
from data.data_types import Video
from tqdm import tqdm


def preload_data() -> Dict[str, Video]:
    """
    Preload data including gallery videos and their posters.
    """
    # Dictionaries for videos and datasets on the backend.
    # Note that since Python 3.7, dictionaries preserve their insert order, so
    # when looping over its `.values()`, elements inserted first also appear first.
    # https://stackoverflow.com/questions/39980323/are-dictionaries-ordered-in-python-3-6
    all_videos = {}

    video_path_pattern = os.path.join(GALLERY_PATH, "**/*.mp4")
    video_paths = glob(video_path_pattern, recursive=True)

    for p in tqdm(video_paths):
        video = get_video(p, GALLERY_PATH)
        all_videos[video.code] = video

    return all_videos


def get_video(
    filepath: os.PathLike,
    absolute_path: Path,
    file_key: Optional[str] = None,
    generate_poster: bool = True,
    width: Optional[int] = None,
    height: Optional[int] = None,
    verbose: Optional[bool] = False,
) -> Video:
    """
    Get video object given
    """
    # Use absolute_path to include the parent directory in the video
    video_path = os.path.relpath(filepath, absolute_path.parent)
    poster_path = None
    if generate_poster:
        poster_id = os.path.splitext(os.path.basename(filepath))[0]
        poster_filename = f"{str(poster_id)}.jpg"
        poster_path = f"{POSTERS_PREFIX}/{poster_filename}"

        # Extract the first frame from video
        poster_output_path = os.path.join(POSTERS_PATH, poster_filename)
        ffmpeg = shutil.which("ffmpeg")
        subprocess.call(
            [
                ffmpeg,
                "-y",
                "-i",
                str(filepath),
                "-pix_fmt",
                "yuv420p",
                "-frames:v",
                "1",
                "-update",
                "1",
                "-strict",
                "unofficial",
                str(poster_output_path),
            ],
            stdout=None if verbose else subprocess.DEVNULL,
            stderr=None if verbose else subprocess.DEVNULL,
        )

        # Extract video width and height from poster. This is important to optimize
        # rendering previews in the mosaic video preview.
        width, height = imagesize.get(poster_output_path)

    return Video(
        code=video_path,
        path=video_path if file_key is None else file_key,
        poster_path=poster_path,
        width=width,
        height=height,
    )