1"""Video loader for multi-modal RAG.
2
3This module provides loading capabilities for video documents with support for:
4- Multiple formats: MP4, AVI, MOV, MKV, FLV, WMV, WebM
5- Key frame extraction
6- Audio track extraction
7- Video metadata extraction
8- Optional transcription of audio track
9"""
10
11from __future__ import annotations
12
13import asyncio
14from pathlib import Path
15import subprocess
16from typing import Any
17
18from lexigram.logging import (
19 get_logger,
20)
21
22logger = get_logger(__name__)
23
24try:
25 import cv2
26 import numpy as np
27
28 CV2_AVAILABLE = True
29except ImportError:
30 cv2 = None
31 np = None # type: ignore[assignment]
32 CV2_AVAILABLE = False
33
34from lexigram.ai.rag.exceptions import VideoLoaderError
35from lexigram.ai.rag.multimodal.loaders.audio import AudioLoader
36from lexigram.ai.rag.multimodal.loaders.image import ImageLoader
37from lexigram.ai.rag.multimodal.types import (
38 AudioDocument,
39 ImageDocument,
40 VideoDocument,
41 VideoFormat,
42 VideoMetadata,
43)
44
45
46class VideoLoader:
47 """Loader for video documents.
48
49 Supports loading videos from files with:
50 - Key frame extraction
51 - Audio track extraction
52 - Metadata extraction
53 - Optional transcription
54
55 Args:
56 num_frames: Number of key frames to extract
57 extract_audio: Whether to extract audio track
58 transcribe_audio: Whether to transcribe audio
59 frame_resize: Maximum dimension for extracted frames
60
61 Example:
62 >>> loader = VideoLoader(num_frames=10, extract_audio=True)
63 >>> doc = await loader.load("/path/to/video.mp4")
64 >>> print(len(doc.frames))
65 10
66 """
67
68 def __init__(
69 self,
70 num_frames: int = 10,
71 extract_audio: bool = True,
72 transcribe_audio: bool = False,
73 frame_resize: int | None = None,
74 ):
75 """Initialize video loader."""
76 if not CV2_AVAILABLE:
77 raise ImportError(
78 "opencv-python is required for video loading. Install with: pip install opencv-python",
79 )
80
81 self.num_frames = num_frames
82 self.extract_audio = extract_audio
83 self.transcribe_audio = transcribe_audio
84 self.frame_resize = frame_resize
85
86 # Initialize sub-loaders
87 self.image_loader = ImageLoader(
88 extract_exif=False,
89 extract_text=False,
90 max_size=frame_resize,
91 )
92
93 if extract_audio:
94 self.audio_loader = AudioLoader(
95 extract_metadata=True,
96 transcribe=transcribe_audio,
97 )
98
99 async def load(
100 self,
101 source: str | Path,
102 title: str | None = None,
103 description: str | None = None,
104 **metadata_kwargs: Any,
105 ) -> VideoDocument:
106 """Load a video document.
107
108 Args:
109 source: File path to video
110 title: Optional video title
111 description: Optional description
112 **metadata_kwargs: Additional metadata fields
113
114 Returns:
115 VideoDocument with loaded video data
116
117 Raises:
118 VideoLoaderError: If loading fails
119 """
120 try:
121 # Validate file exists
122 if isinstance(source, str | Path):
123 file_path = Path(source)
124 exists = await asyncio.to_thread(file_path.exists)
125 if not exists:
126 msg = f"Video file not found: {source}"
127 raise VideoLoaderError(msg)
128 else:
129 raise VideoLoaderError(
130 "Only file paths are supported for video loading",
131 )
132
133 # Detect format (fast)
134 video_format = self._detect_format(file_path)
135
136 # Open video (I/O intensive)
137 def _open_video() -> Any:
138 return cv2.VideoCapture(str(file_path))
139
140 cap = await asyncio.to_thread(_open_video)
141 if not cap.isOpened():
142 msg = f"Failed to open video file: {file_path}"
143 raise VideoLoaderError(msg)
144
145 # Get video properties (fast but safest in thread)
146 def _get_props() -> Any:
147 return {
148 "fps": cap.get(cv2.CAP_PROP_FPS),
149 "frame_count": int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
150 "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
151 "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
152 }
153
154 props = await asyncio.to_thread(_get_props)
155 fps = props["fps"]
156 frame_count = props["frame_count"]
157 width = props["width"]
158 height = props["height"]
159 duration = frame_count / fps if fps > 0 else 0
160
161 # Extract key frames
162 frames = await self._extract_frames(cap, frame_count)
163
164 # Release video
165 await asyncio.to_thread(cap.release)
166
167 # Extract audio track
168 audio_track = None
169 transcript = None
170 if self.extract_audio:
171 try:
172 audio_track = await self._extract_audio(file_path)
173 if audio_track and audio_track.transcript:
174 transcript = audio_track.transcript
175 except (
176 subprocess.SubprocessError,
177 FileNotFoundError,
178 OSError,
179 TimeoutError,
180 ) as e:
181 logger.warning("Failed to extract audio: %s", e)
182
183 # Build metadata
184 metadata = VideoMetadata(
185 title=title,
186 description=description,
187 **metadata_kwargs,
188 )
189
190 # Create document
191 return VideoDocument(
192 content=str(file_path),
193 format=video_format,
194 duration=duration,
195 fps=fps,
196 width=width,
197 height=height,
198 metadata=metadata,
199 frames=frames,
200 audio_track=audio_track,
201 transcript=transcript,
202 file_path=file_path,
203 )
204
205 except Exception as e:
206 msg = f"Failed to load video {source}: {e}"
207 raise VideoLoaderError(msg) from e
208
209 async def load_batch(
210 self,
211 sources: list[str | Path],
212 **kwargs: Any,
213 ) -> list[VideoDocument]:
214 """Load multiple video files.
215
216 Args:
217 sources: List of video sources
218 **kwargs: Common metadata for all videos
219
220 Returns:
221 List of loaded VideoDocuments
222 """
223 if not sources:
224 return []
225
226 async def _safe_load(source) -> Any:
227 try:
228 return await self.load(source, **kwargs)
229 except VideoLoaderError as e:
230 logger.warning("Failed to load video %s: %s", source, e)
231 return None
232
233 tasks = [_safe_load(source) for source in sources]
234 results = await asyncio.gather(*tasks)
235 return [doc for doc in results if doc is not None]
236
237 def _detect_format(self, file_path: Path) -> VideoFormat:
238 """Detect video format from file extension.
239
240 Args:
241 file_path: Path to video file
242
243 Returns:
244 VideoFormat enum value
245 """
246 suffix = file_path.suffix.lower().lstrip(".")
247
248 format_mapping = {
249 "mp4": VideoFormat.MP4,
250 "avi": VideoFormat.AVI,
251 "mov": VideoFormat.MOV,
252 "mkv": VideoFormat.MKV,
253 "flv": VideoFormat.FLV,
254 "wmv": VideoFormat.WMV,
255 "webm": VideoFormat.WEBM,
256 }
257
258 if suffix not in format_mapping:
259 msg = f"Unsupported video format: {suffix}"
260 raise VideoLoaderError(msg)
261
262 return format_mapping[suffix]
263
264 async def _extract_frames(
265 self,
266 cap: cv2.VideoCapture,
267 frame_count: int,
268 ) -> list[ImageDocument]:
269 """Extract key frames from video.
270
271 Args:
272 cap: OpenCV video capture object
273 frame_count: Total number of frames
274
275 Returns:
276 List of ImageDocuments representing key frames
277 """
278 frames: list[ImageDocument] = []
279
280 if frame_count == 0 or self.num_frames == 0:
281 return frames
282
283 # Calculate frame indices to extract (evenly distributed)
284 interval = max(1, frame_count // self.num_frames)
285 frame_indices = [i * interval for i in range(self.num_frames)]
286
287 for idx in frame_indices:
288 # Seek to frame and read (Blocking I/O)
289 def _read_frame() -> Any:
290 cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
291 return cap.read()
292
293 ret, frame = await asyncio.to_thread(_read_frame)
294
295 if not ret:
296 continue
297
298 # Heavy CPU processing in thread
299 def _process_frame() -> Any:
300 # Convert BGR to RGB
301 _frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
302
303 from PIL import Image
304
305 _pil_image = Image.fromarray(_frame_rgb)
306
307 # Resize if needed
308 if self.frame_resize:
309 # Compatibility shim: PIL 10 uses Image.Resampling; older versions use Image.LANCZOS
310 resampling = getattr(Image, "Resampling", Image)
311 resample_filter = getattr(
312 resampling,
313 "LANCZOS",
314 getattr(Image, "LANCZOS", None),
315 )
316
317 if resample_filter is not None:
318 _pil_image.thumbnail(
319 (self.frame_resize, self.frame_resize),
320 resample_filter,
321 )
322 else:
323 _pil_image.thumbnail((self.frame_resize, self.frame_resize))
324
325 # Convert to bytes
326 import io
327
328 _img_byte_arr = io.BytesIO()
329 _pil_image.save(_img_byte_arr, format="JPEG")
330 return _img_byte_arr.getvalue()
331
332 img_bytes = await asyncio.to_thread(_process_frame)
333
334 # Create ImageDocument (async call)
335 frame_doc = await self.image_loader.load(
336 img_bytes,
337 caption=f"Frame {idx}",
338 )
339 frames.append(frame_doc)
340
341 return frames
342
343 async def _extract_audio(self, video_path: Path) -> AudioDocument | None:
344 """Extract audio track from video.
345
346 This uses ffmpeg to extract the audio track to a temporary file,
347 then loads it with the AudioLoader.
348
349 Args:
350 video_path: Path to video file
351
352 Returns:
353 AudioDocument or None if no audio track
354 """
355 import subprocess
356 import tempfile
357
358 # Create temporary audio file
359 with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
360 temp_audio_path = Path(temp_audio.name)
361
362 try:
363 # Extract audio with ffmpeg
364 cmd = [
365 "ffmpeg",
366 "-i",
367 str(video_path),
368 "-vn", # No video
369 "-acodec",
370 "pcm_s16le", # PCM 16-bit little-endian
371 "-ar",
372 "16000", # 16kHz sample rate
373 "-ac",
374 "1", # Mono
375 "-y", # Overwrite output
376 str(temp_audio_path),
377 ]
378
379 # cmd is constructed from fixed ffmpeg args and a validated file path (no shell=True)
380 # Run in thread executor as it is blocking
381 result = await asyncio.to_thread(
382 subprocess.run,
383 cmd,
384 check=False,
385 stdout=subprocess.DEVNULL,
386 stderr=subprocess.DEVNULL,
387 timeout=60,
388 )
389
390 if result.returncode != 0:
391 return None
392
393 # Load audio with AudioLoader
394 audio_doc = await self.audio_loader.load(temp_audio_path)
395 except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
396 return None
397 else:
398 return audio_doc
399 finally:
400 # Clean up temp file
401 if temp_audio_path.exists():
402 temp_audio_path.unlink()