wehub-resource-sync
132 lines
3.3 KiB
Python
132 lines
3.3 KiB
Python
# Copyright (C) CVAT.ai Corporation
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import io
|
|
import subprocess
|
|
from collections.abc import Generator
|
|
from fractions import Fraction
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import av.video.reformatter
|
|
import librosa
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from shared.fixtures.init import get_server_image_tag
|
|
|
|
|
|
def generate_image_file(filename="image.png", size=(100, 50), color=(0, 0, 0)):
|
|
f = BytesIO()
|
|
f.name = filename
|
|
image = Image.new("RGB", size=size, color=color)
|
|
image.save(f)
|
|
f.seek(0)
|
|
|
|
return f
|
|
|
|
|
|
def generate_image_files(
|
|
count: int,
|
|
*,
|
|
prefixes: list[str] | None = None,
|
|
filenames: list[str] | None = None,
|
|
sizes: list[tuple[int, int]] | None = None,
|
|
) -> list[BytesIO]:
|
|
assert not (prefixes and filenames), "prefixes cannot be used together with filenames"
|
|
assert not prefixes or len(prefixes) == count
|
|
assert not filenames or len(filenames) == count
|
|
|
|
images = []
|
|
for i in range(count):
|
|
prefix = prefixes[i] if prefixes else ""
|
|
filename = f"{prefix}{i}.jpeg" if not filenames else filenames[i]
|
|
image = generate_image_file(
|
|
filename, color=(i, i, i), **({"size": sizes[i]}) if sizes else {}
|
|
)
|
|
images.append(image)
|
|
|
|
return images
|
|
|
|
|
|
def generate_video_file(
|
|
num_frames: int,
|
|
*,
|
|
size: tuple[int, int] = (100, 50),
|
|
invalid_keyframes: bool = False,
|
|
) -> BytesIO:
|
|
f = BytesIO()
|
|
f.name = "video.mkv"
|
|
chapters = [
|
|
{
|
|
"id": 0,
|
|
"start": 0,
|
|
"end": 100,
|
|
"time_base": Fraction(1, 1000),
|
|
"metadata": {"title": "Intro"},
|
|
}
|
|
]
|
|
|
|
with av.open(f, "w") as container:
|
|
container.set_chapters(chapters)
|
|
stream = container.add_stream("mjpeg", rate=60)
|
|
stream.width = size[0]
|
|
stream.height = size[1]
|
|
stream.color_range = av.video.reformatter.ColorRange.JPEG
|
|
|
|
for i in range(num_frames):
|
|
frame = av.VideoFrame.from_image(Image.new("RGB", size=size, color=(i, i, i)))
|
|
for packet in stream.encode(frame):
|
|
if invalid_keyframes:
|
|
# Specify pts/dts values that result in 0 valid keyframes
|
|
packet.pts = 10
|
|
packet.dts = 10
|
|
|
|
container.mux(packet)
|
|
|
|
f.seek(0)
|
|
|
|
return f
|
|
|
|
|
|
def read_video_file(file: BytesIO) -> Generator[Image.Image, None, None]:
|
|
file.seek(0)
|
|
|
|
with av.open(file) as container:
|
|
video_stream = container.streams.video[0]
|
|
|
|
for frame in container.decode(video_stream):
|
|
yield frame.to_image()
|
|
|
|
|
|
def generate_manifest(path: str) -> None:
|
|
command = [
|
|
"docker",
|
|
"run",
|
|
"--rm",
|
|
"-u",
|
|
"root:root",
|
|
"-v",
|
|
f"{path}:/local",
|
|
"--entrypoint",
|
|
"python3",
|
|
get_server_image_tag(),
|
|
"utils/dataset_manifest/create.py",
|
|
"--output-dir",
|
|
"/local",
|
|
"/local",
|
|
]
|
|
try:
|
|
subprocess.check_output(command, stderr=subprocess.PIPE)
|
|
except subprocess.CalledProcessError as e:
|
|
print(e.stderr.decode("utf-8"))
|
|
raise
|
|
|
|
|
|
def read_audio_pcm(
|
|
f: Path | io.IOBase, *, offset_ms: int = 0, rate: int = 8000
|
|
) -> tuple[np.ndarray, float]:
|
|
return librosa.load(f, mono=True, sr=rate, offset=offset_ms / 1000)
|