|
from typing import Dict, List, Optional, Tuple, Union |
|
|
|
import cv2 |
|
from omagent_core.utils.logger import logging |
|
from PIL import Image |
|
from pydantic import BaseModel |
|
from pydub import AudioSegment |
|
from pydub.effects import normalize |
|
from scenedetect import (ContentDetector, FrameTimecode, SceneManager, |
|
VideoStream, open_video) |
|
|
|
|
|
class Scene(BaseModel): |
|
start: FrameTimecode |
|
end: FrameTimecode |
|
stt_res: Optional[Dict] = None |
|
summary: Optional[Dict] = None |
|
|
|
class Config: |
|
"""Configuration for this pydantic object.""" |
|
|
|
arbitrary_types_allowed = True |
|
|
|
@classmethod |
|
def init(cls, start: FrameTimecode, end: FrameTimecode, summary: dict = None): |
|
return cls(start=start, end=end, summary=summary) |
|
|
|
@property |
|
def conversation(self): |
|
|
|
if isinstance(self.stt_res, list): |
|
output_conversation = "\n".join( |
|
[f"{item.get('text', None)}" for item in self.stt_res] |
|
) |
|
else: |
|
output_conversation = self.stt_res |
|
return output_conversation |
|
|
|
|
|
class VideoScenes(BaseModel): |
|
stream: VideoStream |
|
audio: Union[AudioSegment, None] |
|
scenes: List[Scene] |
|
frame_extraction_interval: int |
|
|
|
class Config: |
|
"""Configuration for this pydantic object.""" |
|
|
|
extra = "allow" |
|
arbitrary_types_allowed = True |
|
|
|
@classmethod |
|
def load( |
|
cls, |
|
video_path: str, |
|
threshold: int = 27, |
|
min_scene_len: int = 1, |
|
frame_extraction_interval: int = 5, |
|
show_progress: bool = False, |
|
): |
|
"""Load a video file. |
|
|
|
Args: |
|
video_path (str): The path of the video file. Only support local file. |
|
threshold (int): The scene detection threshold. |
|
min_scene_len (int): Once a cut is detected, this long time must pass before a new one can |
|
be added to the scene list. Count in seconds, defaults to 1. |
|
show_progress (bool, optional): Whether to display the progress bar when processing the video. Defaults to False. |
|
""" |
|
video = open_video(video_path) |
|
scene_manager = SceneManager() |
|
scene_manager.add_detector( |
|
ContentDetector( |
|
threshold=threshold, min_scene_len=video.frame_rate * min_scene_len |
|
) |
|
) |
|
scene_manager.detect_scenes(video, show_progress=show_progress) |
|
scenes = scene_manager.get_scene_list(start_in_scene=True) |
|
|
|
try: |
|
audio = AudioSegment.from_file(video_path) |
|
audio = normalize(audio) |
|
except Exception as e: |
|
logging.warning(f"Failed to load audio from {video_path}: {e}") |
|
audio = None |
|
return cls( |
|
stream=video, |
|
scenes=[Scene.init(*scene) for scene in scenes], |
|
audio=audio, |
|
frame_extraction_interval=frame_extraction_interval, |
|
) |
|
|
|
def get_video_frames( |
|
self, scene: Union[int, Scene, Tuple[FrameTimecode]], interval: int = None |
|
) -> Tuple[List[Image.Image], List[float]]: |
|
"""Get the frames of a scene. |
|
|
|
Args: |
|
scene (Union[int, Scene, Tuple[FrameTimecode]]): The scene to get frames. Can be the index of the scene, the scene object or a tuple of start and end frame timecode. |
|
interval (int, optional): The interval of the frames to get. Defaults to None. |
|
Raises: |
|
ValueError: If the type of scene is not int, Scene or tuple. |
|
|
|
Returns: |
|
List[ndarray]: The frames of the scene. |
|
""" |
|
if isinstance(scene, int): |
|
scene = self.scenes[scene] |
|
start, end = scene.start, scene.end |
|
elif isinstance(scene, Scene): |
|
start, end = scene.start, scene.end |
|
elif isinstance(scene, tuple): |
|
start, end = scene |
|
else: |
|
raise ValueError( |
|
f"scene should be int, Scene or tuple, not {type(scene).__name__}" |
|
) |
|
self.stream.seek(start) |
|
frames = [] |
|
time_stamps = [] |
|
if interval is None: |
|
interval = self.frame_extraction_interval * self.stream.frame_rate |
|
scene_len = end.get_frames() - start.get_frames() |
|
if scene_len / 10 > interval: |
|
interval = int(scene_len / 10) + 1 |
|
for index in range(scene_len): |
|
if index % interval == 0: |
|
f = self.stream.read() |
|
if f is False: |
|
continue |
|
f = cv2.cvtColor(f, cv2.COLOR_BGR2RGB) |
|
frames.append(Image.fromarray(f)) |
|
time_stamps.append(self.stream.position.get_seconds()) |
|
else: |
|
self.stream.read(decode=False) |
|
self.stream.seek(0) |
|
return frames, time_stamps |
|
|
|
def get_audio_clip( |
|
self, scene: Union[int, Scene, Tuple[FrameTimecode]] |
|
) -> AudioSegment: |
|
"""Get the audio clip of a scene. |
|
|
|
Args: |
|
scene (Union[int, Scene, Tuple[FrameTimecode]]): The scene to get audio clip. Can be the index of the scene, the scene object or a tuple of start and end frame timecode. |
|
|
|
Raises: |
|
ValueError: If the type of scene is not int, Scene or tuple. |
|
|
|
Returns: |
|
AudioSegment: The audio clip of the scene. |
|
""" |
|
if self.audio is None: |
|
return None |
|
if isinstance(scene, int): |
|
scene = self.scenes[scene] |
|
start, end = scene.start, scene.end |
|
elif isinstance(scene, Scene): |
|
start, end = scene.start, scene.end |
|
elif isinstance(scene, tuple): |
|
start, end = scene |
|
else: |
|
raise ValueError( |
|
f"scene should be int, Scene or tuple, not {type(scene).__name__}" |
|
) |
|
|
|
return self.audio[ |
|
int(start.get_seconds() * 1000) : int(end.get_seconds() * 1000) |
|
] |
|
|
|
def __len__(self): |
|
return len(self.scenes) |
|
|
|
def __iter__(self): |
|
self.index = 0 |
|
return self |
|
|
|
def __next__(self): |
|
if self.index >= len(self.scenes): |
|
raise StopIteration |
|
scene = self.scenes[self.index] |
|
self.index += 1 |
|
return scene |
|
|
|
def __getitem__(self, index): |
|
return self.scenes[index] |
|
|
|
def __setitem__(self, index, value): |
|
self.scenes[index] = value |
|
|
