Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/multimodal/loaders/video.py: 18%

136 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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()