94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
SAMPLE_DIR = ROOT / "storage" / "samples"
|
|
|
|
|
|
def draw_curve(frame: np.ndarray, t: int) -> None:
|
|
height, width = frame.shape[:2]
|
|
points = []
|
|
for x in np.linspace(48, width - 54, 110):
|
|
y = height * 0.55 + 48 * np.sin((x / width) * 2.8 * np.pi + t * 0.11)
|
|
y += 22 * np.sin((x / width) * 7.2 * np.pi + t * 0.035)
|
|
points.append([int(x), int(y)])
|
|
pts = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
|
|
cv2.polylines(frame, [pts], False, (36, 36, 36), 2, cv2.LINE_AA)
|
|
cv2.circle(frame, tuple(points[-1]), 5, (24, 24, 24), -1, cv2.LINE_AA)
|
|
|
|
|
|
def make_frame(index: int, width: int = 640, height: int = 420) -> np.ndarray:
|
|
rng = np.random.default_rng(index)
|
|
base = np.full((height, width), 156, dtype=np.uint8)
|
|
gradient = np.linspace(-22, 18, width, dtype=np.float32)
|
|
base = np.clip(base.astype(np.float32) + gradient[None, :], 0, 255).astype(np.uint8)
|
|
noise = rng.normal(0, 7, (height, width)).astype(np.float32)
|
|
base = np.clip(base.astype(np.float32) + noise, 0, 255).astype(np.uint8)
|
|
frame = cv2.cvtColor(base, cv2.COLOR_GRAY2BGR)
|
|
|
|
for offset in range(-160, width, 130):
|
|
center = (offset + index * 2, height // 2)
|
|
cv2.ellipse(frame, center, (90, 180), 8, 0, 360, (176, 176, 176), 8, cv2.LINE_AA)
|
|
|
|
cv2.line(frame, (80, 60), (560, 86 + index % 18), (90, 90, 90), 4, cv2.LINE_AA)
|
|
draw_curve(frame, index)
|
|
frame = cv2.GaussianBlur(frame, (3, 3), 0)
|
|
return frame
|
|
|
|
|
|
def main() -> None:
|
|
SAMPLE_DIR.mkdir(parents=True, exist_ok=True)
|
|
video_path = SAMPLE_DIR / "synthetic_guidewire.mp4"
|
|
raw_video_path = SAMPLE_DIR / "synthetic_guidewire.raw.mp4"
|
|
image_path = SAMPLE_DIR / "synthetic_guidewire.png"
|
|
width, height = 640, 420
|
|
writer = cv2.VideoWriter(str(raw_video_path), cv2.VideoWriter_fourcc(*"mp4v"), 12.0, (width, height))
|
|
for index in range(72):
|
|
frame = make_frame(index, width, height)
|
|
if index == 12:
|
|
cv2.imwrite(str(image_path), frame)
|
|
writer.write(frame)
|
|
writer.release()
|
|
|
|
ffmpeg = shutil.which("ffmpeg")
|
|
if ffmpeg:
|
|
subprocess.run(
|
|
[
|
|
ffmpeg,
|
|
"-y",
|
|
"-i",
|
|
str(raw_video_path),
|
|
"-c:v",
|
|
"libx264",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-movflags",
|
|
"+faststart",
|
|
"-preset",
|
|
"veryfast",
|
|
"-crf",
|
|
"20",
|
|
str(video_path),
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
raw_video_path.unlink(missing_ok=True)
|
|
else:
|
|
raw_video_path.replace(video_path)
|
|
|
|
print(video_path)
|
|
print(image_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|