225 lines
7.7 KiB
Python
225 lines
7.7 KiB
Python
import subprocess
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
from .ffmpeg_utils import check_ffmpeg
|
|
from .file_utils import resolve_path
|
|
|
|
|
|
TARGET_FPS = 25
|
|
TARGET_AUDIO_SR = 16000
|
|
|
|
|
|
def _as_image_tensor(images):
|
|
if isinstance(images, list):
|
|
images = torch.stack(images)
|
|
if not isinstance(images, torch.Tensor):
|
|
images = torch.as_tensor(images)
|
|
if images.dim() != 4:
|
|
raise ValueError(f"Expected IMAGE tensor with shape [frames, height, width, channels], got {tuple(images.shape)}")
|
|
if images.shape[-1] < 3:
|
|
raise ValueError(f"Expected IMAGE tensor with at least 3 channels, got {tuple(images.shape)}")
|
|
return images.detach().cpu()
|
|
|
|
|
|
def images_to_uint8_rgb(images):
|
|
image_tensor = _as_image_tensor(images)[..., :3]
|
|
if image_tensor.dtype != torch.uint8:
|
|
image_tensor = image_tensor.float()
|
|
if image_tensor.numel() > 0 and float(image_tensor.max()) <= 1.0:
|
|
image_tensor = image_tensor * 255.0
|
|
image_tensor = image_tensor.clamp(0, 255).byte()
|
|
return np.ascontiguousarray(image_tensor.numpy())
|
|
|
|
|
|
def frames_uint8_rgb_to_image_tensor(frames):
|
|
if len(frames) == 0:
|
|
raise ValueError("No output frames were generated.")
|
|
array = np.ascontiguousarray(np.stack(frames).astype(np.float32) / 255.0)
|
|
return torch.from_numpy(array)
|
|
|
|
|
|
def resample_frames_to_fps(frames_rgb, input_fps, output_fps=TARGET_FPS, max_frames=0):
|
|
input_fps = int(input_fps)
|
|
output_fps = int(output_fps)
|
|
if input_fps <= 0 or output_fps <= 0:
|
|
raise ValueError("input_fps and output_fps must be positive.")
|
|
if len(frames_rgb) == 0:
|
|
raise ValueError("No IMAGE frames were provided.")
|
|
|
|
if input_fps == output_fps:
|
|
sampled = frames_rgb.copy()
|
|
else:
|
|
duration = len(frames_rgb) / float(input_fps)
|
|
target_count = max(1, int(round(duration * output_fps)))
|
|
indices = np.floor(np.arange(target_count) * input_fps / output_fps).astype(np.int64)
|
|
indices = np.clip(indices, 0, len(frames_rgb) - 1)
|
|
sampled = frames_rgb[indices]
|
|
|
|
max_frames = int(max_frames or 0)
|
|
if max_frames > 0:
|
|
sampled = sampled[:max_frames]
|
|
|
|
return np.ascontiguousarray(sampled)
|
|
|
|
|
|
def _audio_waveform_to_tensor(audio):
|
|
if not isinstance(audio, dict) or "waveform" not in audio or "sample_rate" not in audio:
|
|
raise ValueError("AUDIO input must be a dict with 'waveform' and 'sample_rate'.")
|
|
|
|
waveform = audio["waveform"]
|
|
if not isinstance(waveform, torch.Tensor):
|
|
waveform = torch.as_tensor(waveform)
|
|
waveform = waveform.detach().cpu().float()
|
|
|
|
if waveform.dim() == 3:
|
|
waveform = waveform[0]
|
|
elif waveform.dim() == 1:
|
|
waveform = waveform.unsqueeze(0)
|
|
elif waveform.dim() != 2:
|
|
raise ValueError(
|
|
"Expected AUDIO waveform with shape [batch, channels, samples], "
|
|
f"[channels, samples], or [samples], got {tuple(waveform.shape)}"
|
|
)
|
|
|
|
return waveform.contiguous(), int(audio["sample_rate"])
|
|
|
|
|
|
def trim_or_pad_audio(audio, duration_sec=None, pad_to_duration=False):
|
|
waveform, sample_rate = _audio_waveform_to_tensor(audio)
|
|
|
|
if duration_sec is not None and duration_sec > 0:
|
|
target_samples = max(1, int(round(float(duration_sec) * sample_rate)))
|
|
if waveform.shape[-1] > target_samples:
|
|
waveform = waveform[:, :target_samples]
|
|
elif pad_to_duration and waveform.shape[-1] < target_samples:
|
|
pad = torch.zeros((waveform.shape[0], target_samples - waveform.shape[-1]), dtype=waveform.dtype)
|
|
waveform = torch.cat([waveform, pad], dim=-1)
|
|
|
|
return {"waveform": waveform.unsqueeze(0).contiguous(), "sample_rate": sample_rate}
|
|
|
|
|
|
def _resample_waveform(waveform, sample_rate, target_sample_rate):
|
|
if int(sample_rate) == int(target_sample_rate):
|
|
return waveform
|
|
|
|
try:
|
|
import torchaudio
|
|
|
|
resampler = torchaudio.transforms.Resample(
|
|
orig_freq=int(sample_rate),
|
|
new_freq=int(target_sample_rate),
|
|
)
|
|
return resampler(waveform)
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
"Failed to resample audio with torchaudio. Install torchaudio or provide 16kHz audio."
|
|
) from exc
|
|
|
|
|
|
def save_audio_to_wav(audio, output_path, duration_sec=None, pad_to_duration=False, mono=False, sample_rate=None):
|
|
adjusted = trim_or_pad_audio(audio, duration_sec=duration_sec, pad_to_duration=pad_to_duration)
|
|
waveform = adjusted["waveform"].squeeze(0).clamp(-1.0, 1.0)
|
|
source_sample_rate = int(adjusted["sample_rate"])
|
|
|
|
if mono and waveform.shape[0] > 1:
|
|
waveform = waveform.mean(dim=0, keepdim=True)
|
|
|
|
target_sample_rate = int(sample_rate or source_sample_rate)
|
|
waveform = _resample_waveform(waveform, source_sample_rate, target_sample_rate)
|
|
|
|
output_file = resolve_path(output_path)
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
audio_np = waveform.transpose(0, 1).numpy()
|
|
sf.write(str(output_file), audio_np, target_sample_rate, subtype="PCM_16")
|
|
return output_file
|
|
|
|
|
|
def encode_rgb_frames_to_mp4(frames_rgb, output_path, fps=TARGET_FPS):
|
|
output_file = resolve_path(output_path)
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
frames = np.ascontiguousarray(frames_rgb)
|
|
if frames.ndim != 4 or frames.shape[-1] != 3:
|
|
raise ValueError(f"Expected RGB frames with shape [frames, height, width, 3], got {frames.shape}")
|
|
|
|
frame_count, height, width, _ = frames.shape
|
|
if frame_count <= 0:
|
|
raise ValueError("No frames were provided to encode.")
|
|
|
|
command = [
|
|
check_ffmpeg(),
|
|
"-y",
|
|
"-f",
|
|
"rawvideo",
|
|
"-vcodec",
|
|
"rawvideo",
|
|
"-s",
|
|
f"{width}x{height}",
|
|
"-pix_fmt",
|
|
"rgb24",
|
|
"-r",
|
|
str(int(fps)),
|
|
"-i",
|
|
"-",
|
|
"-an",
|
|
"-c:v",
|
|
"libx264",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-crf",
|
|
"18",
|
|
str(output_file),
|
|
]
|
|
|
|
process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
try:
|
|
for frame in frames:
|
|
process.stdin.write(np.ascontiguousarray(frame).tobytes())
|
|
_, stderr = process.communicate()
|
|
except BrokenPipeError:
|
|
_, stderr = process.communicate()
|
|
finally:
|
|
if process.stdin and not process.stdin.closed:
|
|
process.stdin.close()
|
|
|
|
if process.returncode != 0:
|
|
message = stderr.decode("utf-8", errors="replace").strip()
|
|
raise RuntimeError(f"Failed to encode HighSync output video: {message}")
|
|
|
|
return output_file
|
|
|
|
|
|
def denoise_audio_file(input_audio_path, output_audio_path, device, cache):
|
|
try:
|
|
from denoiser import pretrained
|
|
from denoiser.dsp import convert_audio
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"denoiser and torchaudio are required when denoise_audio=True. "
|
|
"Install this wrapper's requirements.txt or disable denoise_audio."
|
|
) from exc
|
|
|
|
model = cache.get("denoiser_model")
|
|
model_device = torch.device(device)
|
|
if model is None:
|
|
model = pretrained.dns64().to(model_device)
|
|
model.eval()
|
|
cache["denoiser_model"] = model
|
|
|
|
audio_np, sample_rate = sf.read(str(Path(input_audio_path)), dtype="float32", always_2d=True)
|
|
wav = torch.from_numpy(audio_np.T.copy())
|
|
wav = convert_audio(wav.to(model_device), sample_rate, model.sample_rate, model.chin)
|
|
|
|
with torch.no_grad():
|
|
denoised = model(wav[None])[0]
|
|
|
|
output_file = resolve_path(output_audio_path)
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
audio_np = denoised.detach().cpu().clamp(-1.0, 1.0).numpy()[0]
|
|
sf.write(str(output_file), audio_np, int(model.sample_rate), subtype="PCM_16")
|
|
return output_file
|