wehub-resource-sync
522 lines
20 KiB
Python
522 lines
20 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import base64
|
|
import math
|
|
import numpy as np
|
|
import os
|
|
import re
|
|
import requests
|
|
import torch
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
from requests.adapters import HTTPAdapter
|
|
from typing import Any, Callable, List, TypeVar, Union
|
|
from urllib3.util.retry import Retry
|
|
|
|
from swift.utils import get_env_args
|
|
|
|
# >>> internvl
|
|
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
|
IMAGENET_STD = (0.229, 0.224, 0.225)
|
|
|
|
|
|
def _build_transform(input_size):
|
|
import torchvision.transforms as T
|
|
from torchvision.transforms.functional import InterpolationMode
|
|
MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
|
|
transform = T.Compose([
|
|
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
|
|
T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=MEAN, std=STD)
|
|
])
|
|
return transform
|
|
|
|
|
|
def _find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
|
|
best_ratio_diff = float('inf')
|
|
best_ratio = (1, 1)
|
|
area = width * height
|
|
for ratio in target_ratios:
|
|
target_aspect_ratio = ratio[0] / ratio[1]
|
|
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
|
|
if ratio_diff < best_ratio_diff:
|
|
best_ratio_diff = ratio_diff
|
|
best_ratio = ratio
|
|
elif ratio_diff == best_ratio_diff:
|
|
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
|
|
best_ratio = ratio
|
|
return best_ratio
|
|
|
|
|
|
def _dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
|
|
orig_width, orig_height = image.size
|
|
aspect_ratio = orig_width / orig_height
|
|
|
|
# calculate the existing image aspect ratio
|
|
target_ratios = set((i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1)
|
|
if min_num <= i * j <= max_num)
|
|
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
|
|
|
|
# find the closest aspect ratio to the target
|
|
target_aspect_ratio = _find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
|
|
|
|
# calculate the target width and height
|
|
target_width = image_size * target_aspect_ratio[0]
|
|
target_height = image_size * target_aspect_ratio[1]
|
|
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
|
|
|
|
# resize the image
|
|
resized_img = image.resize((target_width, target_height))
|
|
processed_images = []
|
|
for i in range(blocks):
|
|
box = ((i % (target_width // image_size)) * image_size, (i // (target_width // image_size)) * image_size,
|
|
((i % (target_width // image_size)) + 1) * image_size, ((i //
|
|
(target_width // image_size)) + 1) * image_size)
|
|
# split the image
|
|
split_img = resized_img.crop(box)
|
|
processed_images.append(split_img)
|
|
assert len(processed_images) == blocks
|
|
if use_thumbnail and len(processed_images) != 1:
|
|
thumbnail_img = image.resize((image_size, image_size))
|
|
processed_images.append(thumbnail_img)
|
|
return processed_images
|
|
|
|
|
|
# <<< internvl
|
|
|
|
|
|
def rescale_image(img: Image.Image, max_pixels: int) -> Image.Image:
|
|
import torchvision.transforms as T
|
|
width = img.width
|
|
height = img.height
|
|
if max_pixels is None or max_pixels <= 0 or width * height <= max_pixels:
|
|
return img
|
|
|
|
ratio = width / height
|
|
height_scaled = math.sqrt(max_pixels / ratio)
|
|
width_scaled = height_scaled * ratio
|
|
return T.Resize((int(height_scaled), int(width_scaled)))(img)
|
|
|
|
|
|
_T = TypeVar('_T')
|
|
|
|
|
|
def _check_path(path: str) -> Union[str, None]:
|
|
"""If it is a path, return the string; if it is base64, return None."""
|
|
MAX_PATH_HEURISTIC = 2000
|
|
if len(path) > MAX_PATH_HEURISTIC:
|
|
return
|
|
if os.path.exists(path):
|
|
return os.path.abspath(path)
|
|
data = path
|
|
ROOT_IMAGE_DIR = get_env_args('ROOT_IMAGE_DIR', str, None)
|
|
if ROOT_IMAGE_DIR is not None:
|
|
path = os.path.join(ROOT_IMAGE_DIR, path)
|
|
path = os.path.abspath(os.path.expanduser(path))
|
|
if os.path.exists(path):
|
|
return path
|
|
if data.startswith('data:'):
|
|
return
|
|
try:
|
|
base64.b64decode(data)
|
|
return
|
|
except Exception:
|
|
pass
|
|
return data
|
|
|
|
|
|
def load_file(path: Union[str, bytes, _T]) -> Union[BytesIO, _T]:
|
|
res = path
|
|
if isinstance(path, str):
|
|
path = path.strip()
|
|
if path.startswith('http'):
|
|
retries = Retry(total=3, backoff_factor=1, allowed_methods=['GET'])
|
|
with requests.Session() as session:
|
|
session.mount('http://', HTTPAdapter(max_retries=retries))
|
|
session.mount('https://', HTTPAdapter(max_retries=retries))
|
|
|
|
timeout = float(os.getenv('SWIFT_TIMEOUT', '20'))
|
|
request_kwargs = {'timeout': timeout} if timeout > 0 else {}
|
|
|
|
response = session.get(path, **request_kwargs)
|
|
response.raise_for_status()
|
|
content = response.content
|
|
res = BytesIO(content)
|
|
else:
|
|
data = path
|
|
path = _check_path(path)
|
|
if path is None:
|
|
# base64_str
|
|
if data.startswith('data:'):
|
|
match_ = re.match(r'data:(.+?);base64,(.+)', data)
|
|
assert match_ is not None
|
|
data = match_.group(2)
|
|
data = base64.b64decode(data)
|
|
res = BytesIO(data)
|
|
else:
|
|
with open(path, 'rb') as f:
|
|
res = BytesIO(f.read())
|
|
elif isinstance(path, bytes):
|
|
res = BytesIO(path)
|
|
return res
|
|
|
|
|
|
def load_image(image: Union[str, bytes, Image.Image]) -> Image.Image:
|
|
image = load_file(image)
|
|
if isinstance(image, BytesIO):
|
|
image = Image.open(image)
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
return image
|
|
|
|
|
|
def load_batch(path_list: List[Union[str, None, Any, BytesIO]],
|
|
load_func: Callable[[Any], _T] = load_image) -> List[_T]:
|
|
res = []
|
|
assert isinstance(path_list, (list, tuple)), f'path_list: {path_list}'
|
|
for path in path_list:
|
|
if path is None: # ignore None
|
|
continue
|
|
res.append(load_func(path))
|
|
return res
|
|
|
|
|
|
def load_video_hf(videos: List[str]):
|
|
from transformers.video_utils import load_video
|
|
res = []
|
|
video_metadata = []
|
|
for video in videos:
|
|
if isinstance(video, (list, tuple)) and isinstance(video[0], str):
|
|
# Case a: Video is provided as a list of image file names
|
|
video = [np.array(load_image(image_fname)) for image_fname in video]
|
|
video = np.stack(video)
|
|
metadata = None
|
|
else:
|
|
# Case b: Video is provided as a single file path or URL or decoded frames in a np.ndarray or torch.tensor
|
|
video_load_backend = get_env_args('video_load_backend', str, 'pyav')
|
|
video, metadata = load_video(
|
|
video,
|
|
backend=video_load_backend,
|
|
)
|
|
res.append(video)
|
|
video_metadata.append(metadata)
|
|
return res, video_metadata
|
|
|
|
|
|
def _get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
|
|
if bound:
|
|
start, end = bound[0], bound[1]
|
|
else:
|
|
start, end = -100000, 100000
|
|
start_idx = max(first_idx, round(start * fps))
|
|
end_idx = min(round(end * fps), max_frame)
|
|
seg_size = float(end_idx - start_idx) / num_segments
|
|
frame_indices = np.array(
|
|
[int(start_idx + (seg_size / 2) + np.round(seg_size * idx)) for idx in range(num_segments)])
|
|
return frame_indices
|
|
|
|
|
|
def transform_image(image, input_size=448, max_num=12):
|
|
transform = _build_transform(input_size=input_size)
|
|
images = _dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
|
|
pixel_values = [transform(image) for image in images]
|
|
pixel_values = torch.stack(pixel_values)
|
|
return pixel_values
|
|
|
|
|
|
def load_video_internvl(video: Union[str, bytes], bound=None, num_segments=32):
|
|
from decord import VideoReader, cpu
|
|
video_io = load_file(video)
|
|
vr = VideoReader(video_io, ctx=cpu(0), num_threads=1)
|
|
max_frame = len(vr) - 1
|
|
fps = float(vr.get_avg_fps())
|
|
|
|
images = []
|
|
frame_indices = _get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
|
|
for frame_index in frame_indices:
|
|
images.append(Image.fromarray(vr[frame_index].asnumpy()).convert('RGB'))
|
|
return images
|
|
|
|
|
|
def load_video_cogvlm2(video: Union[str, bytes]) -> np.ndarray:
|
|
from decord import VideoReader, bridge, cpu
|
|
video_io = load_file(video)
|
|
bridge.set_bridge('torch')
|
|
clip_end_sec = 60
|
|
clip_start_sec = 0
|
|
num_frames = get_env_args('num_frames', int, 24)
|
|
decord_vr = VideoReader(video_io, ctx=cpu(0))
|
|
duration = len(decord_vr) # duration in terms of frames
|
|
start_frame = int(clip_start_sec * decord_vr.get_avg_fps())
|
|
end_frame = min(duration, int(clip_end_sec * decord_vr.get_avg_fps())) if \
|
|
clip_end_sec is not None else duration
|
|
frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int)
|
|
video_data = decord_vr.get_batch(frame_id_list)
|
|
video_data = video_data.permute(3, 0, 1, 2)
|
|
return video_data
|
|
|
|
|
|
def load_video_llava(video: Union[str, bytes]) -> np.ndarray:
|
|
import av
|
|
video_io = load_file(video)
|
|
container = av.open(video_io)
|
|
total_frames = container.streams.video[0].frames
|
|
num_frames = get_env_args('num_frames', int, 16)
|
|
indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
|
|
frames = []
|
|
container.seek(0)
|
|
start_index = indices[0]
|
|
end_index = indices[-1]
|
|
for i, frame in enumerate(container.decode(video=0)):
|
|
if i > end_index:
|
|
break
|
|
if i >= start_index and i in indices:
|
|
frames.append(frame)
|
|
return np.stack([x.to_ndarray(format='rgb24') for x in frames])
|
|
|
|
|
|
def load_video_minicpmv_mplug_owl3(video: Union[str, bytes], max_num_frames):
|
|
|
|
from decord import VideoReader, cpu # pip install decord
|
|
|
|
def uniform_sample(_l, _n):
|
|
gap = len(_l) / _n
|
|
idxs = [int(i * gap + gap / 2) for i in range(_n)]
|
|
return [_l[i] for i in idxs]
|
|
|
|
video_io = load_file(video)
|
|
vr = VideoReader(video_io, ctx=cpu(0))
|
|
sample_fps = round(vr.get_avg_fps() / 1) # FPS
|
|
frame_idx = [i for i in range(0, len(vr), sample_fps)]
|
|
|
|
if len(frame_idx) > max_num_frames:
|
|
frame_idx = uniform_sample(frame_idx, max_num_frames)
|
|
frames = vr.get_batch(frame_idx).asnumpy()
|
|
frames = [Image.fromarray(v.astype('uint8')) for v in frames]
|
|
return frames
|
|
|
|
|
|
def _load_audio_librosa(audio: Union[str, bytes], sampling_rate: int, mono: bool = True):
|
|
import librosa
|
|
try:
|
|
audio_io = load_file(audio)
|
|
return librosa.load(audio_io, sr=sampling_rate, mono=mono)
|
|
except Exception:
|
|
if isinstance(audio, str) and audio.startswith(('http://', 'https://')):
|
|
import audioread
|
|
audio_io = audioread.ffdec.FFmpegAudioFile(audio)
|
|
else:
|
|
audio_io = _check_path(audio) if isinstance(audio, str) else audio
|
|
return librosa.load(audio_io, sr=sampling_rate, mono=mono)
|
|
|
|
|
|
# ref: https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/multimodal/audio.py#L169-L224
|
|
def _resample_audio_pyav(audio: np.ndarray, *, orig_sr: float, target_sr: float) -> np.ndarray:
|
|
import av
|
|
orig_sr_int = int(round(orig_sr))
|
|
target_sr_int = int(round(target_sr))
|
|
if orig_sr_int == target_sr_int:
|
|
return audio
|
|
if audio.ndim == 2:
|
|
return np.stack([_resample_audio_pyav(ch, orig_sr=orig_sr, target_sr=target_sr) for ch in audio], axis=0)
|
|
expected_len = int(math.ceil(audio.shape[-1] * target_sr_int / orig_sr_int))
|
|
min_samples = 1024
|
|
audio_f32 = np.asarray(audio, dtype=np.float32)
|
|
if len(audio_f32) < min_samples:
|
|
audio_f32 = np.pad(audio_f32, (0, min_samples - len(audio_f32)))
|
|
audio_f32 = audio_f32.reshape(1, -1)
|
|
resampler = av.AudioResampler(format='fltp', layout='mono', rate=target_sr_int)
|
|
frame = av.AudioFrame.from_ndarray(audio_f32, format='fltp', layout='mono')
|
|
frame.sample_rate = orig_sr_int
|
|
out_frames = resampler.resample(frame)
|
|
out_frames.extend(resampler.resample(None))
|
|
result = np.concatenate([f.to_ndarray() for f in out_frames], axis=1).squeeze(0)
|
|
return result[:expected_len]
|
|
|
|
|
|
# ref: https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/multimodal/media/audio.py#L45-L160
|
|
def _load_audio_soundfile_pyav(path: Union[str, bytes, BytesIO], *, sr: float, mono: bool = True):
|
|
"""soundfile first, pyav fallback — same strategy as vLLM multimodal audio loader."""
|
|
bad_sf_codes = {0, 1, 3, 4}
|
|
if not isinstance(path, BytesIO):
|
|
path = load_file(path)
|
|
|
|
def _load_soundfile():
|
|
import soundfile
|
|
with soundfile.SoundFile(path) as f:
|
|
native_sr = f.samplerate
|
|
y = f.read(dtype='float32', always_2d=False).T
|
|
|
|
if mono and y.ndim > 1:
|
|
y = np.mean(y, axis=tuple(range(y.ndim - 1)))
|
|
|
|
if sr is not None and sr != native_sr:
|
|
y = _resample_audio_pyav(y, orig_sr=native_sr, target_sr=sr)
|
|
return y, int(sr)
|
|
return y, native_sr
|
|
|
|
def _load_pyav():
|
|
import av
|
|
path.seek(0)
|
|
with av.open(path) as container:
|
|
if not container.streams.audio:
|
|
raise ValueError('No audio stream found.')
|
|
stream = container.streams.audio[0]
|
|
stream.thread_type = 'AUTO'
|
|
native_sr = stream.rate
|
|
target_sr = sr or native_sr
|
|
chunks = []
|
|
needs_resampling = not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6)
|
|
resampler = av.AudioResampler(format='fltp', layout='mono', rate=target_sr) if needs_resampling else None
|
|
for frame in container.decode(stream):
|
|
if needs_resampling:
|
|
for out_frame in resampler.resample(frame):
|
|
chunks.append(out_frame.to_ndarray())
|
|
else:
|
|
chunks.append(frame.to_ndarray())
|
|
if not chunks:
|
|
raise ValueError('No audio found.')
|
|
y = np.concatenate(chunks, axis=-1).astype(np.float32)
|
|
if mono and y.ndim > 1:
|
|
y = np.mean(y, axis=0)
|
|
return y, target_sr
|
|
|
|
try:
|
|
return _load_soundfile()
|
|
except ImportError:
|
|
path.seek(0)
|
|
return _load_pyav()
|
|
except Exception as exc:
|
|
import soundfile
|
|
if not isinstance(exc, soundfile.LibsndfileError) or exc.code not in bad_sf_codes:
|
|
raise
|
|
path.seek(0)
|
|
return _load_pyav()
|
|
|
|
|
|
def load_audio(
|
|
audio: Union[str, bytes, BytesIO],
|
|
sampling_rate: int,
|
|
return_sr: bool = False,
|
|
mono: bool = True,
|
|
):
|
|
backend = get_env_args('swift_audio_load_backend', str, 'librosa')
|
|
if backend == 'librosa':
|
|
res = _load_audio_librosa(audio, sampling_rate, mono=mono)
|
|
elif backend == 'soundfile_pyav':
|
|
res = _load_audio_soundfile_pyav(audio, sr=sampling_rate, mono=mono)
|
|
else:
|
|
raise ValueError(f'Unknown audio load backend {backend!r}. Supported: librosa, soundfile_pyav')
|
|
return res if return_sr else res[0]
|
|
|
|
|
|
def _resolve_video_local_path(path: Union[str, bytes]) -> tuple:
|
|
"""Return (local_path, is_temp_file). HTTP URLs and raw bytes are written to a temp file."""
|
|
if isinstance(path, bytes) or (isinstance(path, str) and path.startswith('http')):
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as f:
|
|
f.write(load_file(path).read())
|
|
return f.name, True
|
|
checked = _check_path(path) if isinstance(path, str) else None
|
|
return checked or path, False
|
|
|
|
|
|
def _video_to_ndarrays_local(local_path: str, num_frames: int = -1) -> np.ndarray:
|
|
import cv2
|
|
cap = cv2.VideoCapture(local_path)
|
|
if not cap.isOpened():
|
|
raise ValueError(f'Could not open video file {local_path}')
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
if total_frames <= 0:
|
|
cap.release()
|
|
raise ValueError(f'Video file {local_path} has invalid or zero frame count: {total_frames}')
|
|
if num_frames <= 0 or num_frames > total_frames:
|
|
num_frames = total_frames
|
|
frame_indices = set(np.linspace(0, total_frames - 1, num_frames, dtype=int))
|
|
frames = []
|
|
for idx in range(total_frames):
|
|
if not cap.grab():
|
|
break
|
|
if idx in frame_indices:
|
|
ret, frame = cap.retrieve()
|
|
if ret:
|
|
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
cap.release()
|
|
if len(frames) < num_frames:
|
|
raise ValueError(f'Could not read enough frames from video file {local_path} '
|
|
f'(expected {num_frames} frames, got {len(frames)})')
|
|
return np.stack(frames)
|
|
|
|
|
|
def _video_get_metadata_local(local_path: str, num_frames: int = -1) -> dict:
|
|
import cv2
|
|
cap = cv2.VideoCapture(local_path)
|
|
if not cap.isOpened():
|
|
raise ValueError(f'Could not open video file {local_path}')
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
if total_frames <= 0:
|
|
cap.release()
|
|
raise ValueError(f'Video file {local_path} has invalid or zero frame count: {total_frames}')
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
duration = total_frames / fps if fps > 0 else 0
|
|
cap.release()
|
|
if num_frames <= 0 or num_frames > total_frames:
|
|
num_frames = total_frames
|
|
return {
|
|
'total_num_frames': num_frames,
|
|
'fps': duration / num_frames if num_frames else fps,
|
|
'duration': duration,
|
|
'video_backend': 'opencv',
|
|
'frames_indices': list(range(num_frames)),
|
|
'do_sample_frames': num_frames == total_frames,
|
|
}
|
|
|
|
|
|
def load_vllm_video(path: Union[str, bytes], num_frames: int = -1) -> tuple:
|
|
"""Decode video frames + metadata for vLLM rollout; one download, temp file cleaned up."""
|
|
local_path, is_temp = _resolve_video_local_path(path)
|
|
try:
|
|
return _video_to_ndarrays_local(local_path, num_frames), _video_get_metadata_local(local_path, num_frames)
|
|
finally:
|
|
if is_temp:
|
|
try:
|
|
os.remove(local_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def load_video_valley(video: Union[str, bytes]):
|
|
import decord
|
|
from torchvision import transforms
|
|
video_io = load_file(video)
|
|
video_reader = decord.VideoReader(video_io)
|
|
decord.bridge.set_bridge('torch')
|
|
video = video_reader.get_batch(np.linspace(0, len(video_reader) - 1, 8).astype(np.int_)).byte()
|
|
images = [transforms.ToPILImage()(image.permute(2, 0, 1)).convert('RGB') for image in video]
|
|
return images
|
|
|
|
|
|
def load_video_ovis2(video_path, num_frames):
|
|
from moviepy.editor import VideoFileClip
|
|
with VideoFileClip(video_path) as clip:
|
|
total_frames = int(clip.fps * clip.duration)
|
|
if total_frames <= num_frames:
|
|
sampled_indices = range(total_frames)
|
|
else:
|
|
stride = total_frames / num_frames
|
|
sampled_indices = [
|
|
min(total_frames - 1, int((stride * i + stride * (i + 1)) / 2)) for i in range(num_frames)
|
|
]
|
|
frames = [clip.get_frame(index / clip.fps) for index in sampled_indices]
|
|
frames = [Image.fromarray(frame, mode='RGB') for frame in frames]
|
|
return frames
|
|
|
|
|
|
def load_video_ovis2_5(video_path, num_frames):
|
|
from moviepy.editor import VideoFileClip
|
|
with VideoFileClip(video_path) as clip:
|
|
total_frames = int(clip.fps * clip.duration)
|
|
indices = [int(i * total_frames / num_frames) for i in range(num_frames)]
|
|
frames = [Image.fromarray(clip.get_frame(t)) for t in (idx / clip.fps for idx in indices)]
|
|
return frames
|