61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
from PIL import Image
|
|
import numpy as np
|
|
from io import BytesIO
|
|
import torch
|
|
from imageio_ffmpeg import get_ffmpeg_exe
|
|
import subprocess
|
|
import re
|
|
|
|
FFMPEG_PATH = get_ffmpeg_exe()
|
|
ENCODE_ARGS = ("utf-8", 'backslashreplace')
|
|
|
|
def get_image_mask(image_data):
|
|
try:
|
|
image = Image.open(BytesIO(image_data))
|
|
|
|
if "A" in image.getbands():
|
|
mask = np.array(image.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
|
|
image = image.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
return image, mask
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load Base64 image: {e}")
|
|
|
|
def get_audio(audio_data, start_time=0, duration=0):
|
|
try:
|
|
args = [FFMPEG_PATH, '-i', 'pipe:0', '-vn', '-acodec', 'pcm_f32le', '-f', 'f32le', 'pipe:1']
|
|
|
|
if start_time > 0:
|
|
args.insert(1, '-ss')
|
|
args.insert(2, str(start_time))
|
|
if duration > 0:
|
|
args.insert(1, '-t')
|
|
args.insert(2, str(duration))
|
|
|
|
try:
|
|
res = subprocess.run(args, input=audio_data, capture_output=True, check=True)
|
|
|
|
audio = torch.frombuffer(res.stdout, dtype=torch.float32)
|
|
|
|
match = re.search(r', (\d+) Hz, (\w+),', res.stderr.decode(*ENCODE_ARGS))
|
|
if match:
|
|
sample_rate = int(match.group(1))
|
|
channels = {"mono": 1, "stereo": 2}[match.group(2)]
|
|
else:
|
|
sample_rate = 44100
|
|
channels = 2
|
|
|
|
audio = audio.reshape((-1, channels)).transpose(0, 1).unsqueeze(0)
|
|
|
|
return {'waveform': audio, 'sample_rate': sample_rate}
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
raise RuntimeError(f"FFmpeg failed:\n{e.stderr.decode(*ENCODE_ARGS)}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load Base64 audio: {e}") |