from typing import Dict, List, Optional, Tuple, Union 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): # for self deployed whisper 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, kernel_size: Optional[int] = None, ): """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() weight = ContentDetector.Components( delta_hue=1.0, delta_sat=1.0, delta_lum=0.0, delta_edges=1.0, ) if kernel_size is None: scene_manager.add_detector( ContentDetector( threshold=threshold, min_scene_len=int(video.frame_rate * min_scene_len), weights=weight, ) ) else: scene_manager.add_detector( ContentDetector( threshold=threshold, min_scene_len=int(video.frame_rate * min_scene_len), weights=weight, kernel_size=kernel_size, ) ) 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 (IndexError, OSError): 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() 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 def to_serializable(self) -> dict: """Convert VideoScenes to a serializable dictionary.""" scenes_data = [] for scene in self.scenes: scenes_data.append( { "start_frame": scene.start.frame_num, "end_frame": scene.end.frame_num, "stt_res": scene.stt_res, "summary": scene.summary, } ) return { "video_path": self.stream.path, "frame_rate": self.stream.frame_rate, "scenes": scenes_data, "frame_extraction_interval": self.frame_extraction_interval, } @classmethod def from_serializable(cls, data: dict): """Rebuild VideoScenes from serialized data.""" video = open_video(data["video_path"]) try: audio = AudioSegment.from_file(data["video_path"]) audio = normalize(audio) except Exception: audio = None # Rebuild scenes list scenes = [] for scene_data in data["scenes"]: start = FrameTimecode(scene_data["start_frame"], data["frame_rate"]) end = FrameTimecode(scene_data["end_frame"], data["frame_rate"]) scene = Scene.init(start, end) scene.stt_res = scene_data["stt_res"] scene.summary = scene_data["summary"] scenes.append(scene) return cls( stream=video, scenes=scenes, audio=audio, frame_extraction_interval=data["frame_extraction_interval"], )