Coverage for apps/streamshort/streamshort_job.py: 70%
298 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamShort job to generate a video short.
3It coordinates the execution of the different models.
4"""
6import sys
7import json
8import cv2
9import aiofiles
10import asyncio
12from dataclasses import asdict
14from typing import override
15from typing import Dict
16from typing import Any
17from typing import List
18from typing import Optional
20from scenedetect import open_video
21from scenedetect import SceneManager
22from scenedetect.detectors import ContentDetector
23from scenedetect.stats_manager import StatsManager
25from short_prompts import DESCRIPTION_PROMPT
26from short_prompts import HIGHLIGHT_PROMPT
28# Local relative imports
29sys.path.append("..") # noqa: E402
30sys.path.append("../..") # noqa: E402
32from streamwise_job import StreamWiseJob
33from streamwise_job import JobStatus
35from scene import SceneSegment
37from lmm_service_manager import LMMServiceManager
39from console_utils import bytes_to_human
41from file_utils import save_base64_as_binary
42from file_utils import read_file_base64
43from file_utils import read_file_bytes
44from media_utils import chunk_video_binary
45from media_utils import concatenate_videos
46from media_utils import extract_audio_from_video
47from media_utils import chunk_audio_base64
50MAX_KEY_FRAMES = 64
51DESCRIPTION_BATCH_SIZE = 16
54class StreamShortJob(StreamWiseJob):
55 """A job to generate a short video summary."""
57 def __init__(
58 self,
59 job_id: str,
60 service_manager: LMMServiceManager,
61 config: Dict[str, Any] = {},
62 ) -> None:
63 super().__init__(
64 "streamshort",
65 job_id,
66 service_manager,
67 config)
68 self.scenes: List[SceneSegment] = []
70 @override
71 async def generate(
72 self,
73 job_config: Dict[str, Any],
74 ) -> None:
75 video_base64: Optional[str] = job_config.get("video_base64")
76 await self.gen_short(video_base64)
78 def find_scene_for_frame(
79 self,
80 frame_num: int
81 ) -> Optional[SceneSegment]:
82 """Find the scene segment that contains the given frame number."""
83 scenes = self.scenes
84 left = 0
85 right = len(scenes) - 1
86 while left <= right:
87 mid = (left + right) // 2
88 scene = scenes[mid]
89 if frame_num < scene.start_frame:
90 right = mid - 1
91 elif frame_num >= scene.end_frame:
92 left = mid + 1
93 else:
94 return scene
95 return None
97 async def gen_short(
98 self,
99 video_base64: Optional[str],
100 ) -> None:
101 """
102 Generate the video short.
103 """
104 async with self.job_status_handler():
105 if not video_base64:
106 self.logger.error("Video is required.")
107 await self.save_status(JobStatus.FAILED)
108 raise ValueError("Missing 'video_base64' in request")
109 self.logger.info(f"Generating short for video with {bytes_to_human(len(video_base64))}.")
111 # Save as video for debugging
112 self.logger.info(f"Saving input video with {bytes_to_human(len(video_base64))}.")
113 video_path = f"{self.job_path}/video.mp4"
114 await save_base64_as_binary(video_path, video_base64)
116 await self.save_status(JobStatus.RUNNING)
118 # Detect scenes
119 self.scenes = await self.detect_scenes()
120 self.logger.info(f"Detected {len(self.scenes)} scenes.")
122 await self.save_status(JobStatus.RUNNING)
124 # Extract key frames
125 key_frames = self.extract_key_frames(video_path)
126 self.logger.info(f"Extracted {len(key_frames)} key frames: {', '.join(map(str, key_frames))}.")
128 # Transcribe audio
129 await self.chunk_audio_into_scenes()
131 transcript_task = asyncio.create_task(self.transcribe_audio())
133 # Describe key frames
134 description_task = asyncio.create_task(self.describe_frames(key_frames))
136 await self.save_status(JobStatus.RUNNING)
138 # Wait for async tasks to complete
139 await transcript_task
140 await description_task
142 # Output some debug files
143 self.logger.info("Scenes:")
144 for idx, scene in enumerate(self.scenes):
145 self.logger.info(f" Scene {idx}: {scene}")
146 scenes_path = f"{self.job_path}/scenes.json"
147 async with aiofiles.open(scenes_path, "w") as scene_file:
148 scenes_dict_list = [asdict(scene) for scene in self.scenes]
149 scenes_json = json.dumps(scenes_dict_list, indent=2)
150 await scene_file.write(scenes_json)
152 await self.save_status(JobStatus.RUNNING)
154 # Select highlight scenes
155 avg_scene_seconds = sum(scene.duration_sec for scene in self.scenes) / max(len(self.scenes), 1)
156 video_duration_seconds = self.get_config_float("video_duration_seconds", 10.0)
157 max_scenes = 1
158 if avg_scene_seconds > 0:
159 max_scenes = max(1, int(video_duration_seconds / avg_scene_seconds))
160 self.logger.info(f"Generating a short of {video_duration_seconds} seconds and max {max_scenes} scenes.")
161 chosen_scenes = await self.choose_scenes_for_highlight(
162 total_length=video_duration_seconds,
163 max_scenes=max_scenes,
164 )
165 short_duration_seconds = sum(self.scenes[scene_id].duration_sec for scene_id in chosen_scenes)
166 self.logger.info(
167 f"Chosen {len(chosen_scenes)} scenes for a short of {short_duration_seconds:.1f} seconds: "
168 f"{', '.join(map(str, chosen_scenes))}.")
169 if not chosen_scenes:
170 raise RuntimeError("Failed to choose scenes for highlight short")
172 # Save selected scene IDs so the WebUI can highlight them
173 await self.save_selected_scenes(chosen_scenes)
175 # Generate some video using diffusion
176 # TODO
178 # Build short
179 short_video_path = await self.save_highlight_short(chosen_scenes)
180 self.logger.info(f"Saved short video to {short_video_path}.")
182 def extract_key_frames(
183 self,
184 video_path: str
185 ) -> List[int]:
186 """Extract key frames from the video and associate them with scenes."""
187 cap = cv2.VideoCapture(video_path)
188 key_frames = self.pick_key_frames()
190 for frame_num in key_frames:
191 frame_file_name = f"frame_{frame_num:04d}.jpg"
192 frame_path = f"{self.job_path}/{frame_file_name}"
193 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
194 ok, frame = cap.read()
195 if not ok:
196 raise RuntimeError(f"Failed to read frame {frame_num}")
197 cv2.imwrite(frame_path, frame)
199 scene = self.find_scene_for_frame(frame_num)
200 if scene:
201 scene.add_image_path(frame_file_name)
202 cap.release()
204 return key_frames
206 async def detect_scenes(
207 self,
208 threshold: float = 27.0,
209 min_scene_len: int = 15,
210 ) -> List[SceneSegment]:
211 """Return list of (start_frame, end_frame, start_sec, end_sec)."""
212 video_path = f"{self.job_path}/video.mp4"
213 if not await aiofiles.os.path.exists(video_path):
214 raise FileNotFoundError(f"Video file not found: {video_path}")
215 video = open_video(video_path)
217 stats_manager = StatsManager()
218 scene_manager = SceneManager(stats_manager)
219 content_detector = ContentDetector(
220 threshold=threshold,
221 min_scene_len=min_scene_len)
222 scene_manager.add_detector(content_detector)
224 scene_manager.detect_scenes(video)
225 scene_list = scene_manager.get_scene_list()
227 scenes = []
228 for scene_id, (start_tc, end_tc) in enumerate(scene_list):
229 scene = SceneSegment(
230 scene_id,
231 start_tc.get_frames(), end_tc.get_frames(),
232 start_tc.get_seconds(), end_tc.get_seconds()
233 )
234 scenes.append(scene)
235 return scenes
237 def pick_key_frames(
238 self,
239 max_frames: int = MAX_KEY_FRAMES,
240 ) -> List[int]:
241 """Pick up to max_frames evenly spaced frame numbers within a scene."""
242 if not self.scenes:
243 return []
245 start_frame = 0
246 end_frame = self.scenes[-1].end_frame
248 if max_frames <= 1 or end_frame <= start_frame:
249 return [max(start_frame, 0)]
250 total_frames = max(end_frame - start_frame, 1)
251 count = min(max_frames, total_frames)
252 frames = []
253 for i in range(count):
254 pos = start_frame + int((i + 0.5) * total_frames / count)
255 pos = min(pos, end_frame - 1)
256 frames.append(pos)
258 return sorted(set(frames))
260 async def describe_frames(
261 self,
262 key_frames: List[int]
263 ) -> None:
264 description_tasks = []
266 key_frame_batch = []
267 batch_id = 0
268 for frame_num in key_frames:
269 key_frame_batch.append(frame_num)
270 if len(key_frame_batch) >= DESCRIPTION_BATCH_SIZE:
271 description_task = asyncio.create_task(self.describe_frames_batch(
272 key_frame_batch, batch_id=batch_id))
273 description_tasks.append(description_task)
274 batch_id += 1
275 key_frame_batch = []
276 if key_frame_batch:
277 description_task = asyncio.create_task(self.describe_frames_batch(
278 key_frame_batch,
279 batch_id=batch_id))
280 description_tasks.append(description_task)
282 for description_task in description_tasks:
283 await description_task
285 async def describe_frames_batch(
286 self,
287 frame_nums: List[int],
288 max_tokens: int = 8192,
289 batch_id: int = 0,
290 ) -> None:
291 """Describe a list of frames using the LLM."""
292 MAX_LOG_TEXT = 80
294 self.logger.info(f"Describing {len(frame_nums)} key frames: {', '.join(map(str, frame_nums))}.")
295 image_base64s = []
296 for frame_num in frame_nums:
297 frame_file_name = f"frame_{frame_num:04d}.jpg"
298 frame_path = f"{self.job_path}/{frame_file_name}"
299 frame_base64 = await read_file_base64(frame_path)
300 image_base64s.append(frame_base64)
302 # Send the frames to the LLM
303 content: List[Dict[str, Any]] = [
304 {"type": "text", "text": DESCRIPTION_PROMPT}
305 ]
306 for image_base64 in image_base64s:
307 content.append({
308 "type": "image_url",
309 "image_url": {
310 "url": f"data:image/jpeg;base64,{image_base64}"
311 }
312 })
313 content.append({"type": "text", "text": "Generate the JSON descriptions now."})
314 message = {
315 "role": "user",
316 "content": content,
317 }
318 messages = [message]
320 prompt_path = f"{self.job_path}/description_prompt_{batch_id}.json"
321 async with aiofiles.open(prompt_path, "w") as prompt_file:
322 await prompt_file.write(json.dumps(messages, indent=2))
324 # Query the LLM
325 response_message = await self.gen.gen_text(
326 messages,
327 max_tokens=max_tokens,
328 task_id=f"describe{batch_id:03d}",
329 )
330 """
331 response = await self.llm_client.chat.completions.create(
332 model=self.llm_model,
333 messages=messages,
334 max_tokens=max_tokens,
335 extra_body=None,
336 # stream=True,
337 )
338 response_message = response.choices[0].message
339 response_message = response_message.content.strip()
340 """
342 prompt_path = f"{self.job_path}/description_response_{batch_id}.txt"
343 async with aiofiles.open(prompt_path, "w") as prompt_file:
344 await prompt_file.write(json.dumps(response_message, indent=2))
346 # Parse multiple descriptions from response
347 self.logger.info("Frames:")
348 num_described_frames = 0
349 for response_line in response_message.splitlines():
350 response_line_strip = response_line.strip()
351 response_line_strip = response_line_strip.strip(",")
352 if not response_line_strip:
353 pass
354 elif "```" in response_line_strip or response_line_strip == "[" or response_line_strip == "]":
355 pass
356 elif not response_line_strip.startswith("{"):
357 self.logger.debug(f"Skipping: {response_line_strip}")
358 else:
359 try:
360 response_json = json.loads(response_line_strip)
361 frame_num_response = response_json.get("frame_num", None)
362 description = response_json.get("description", None)
363 if not description:
364 self.logger.warning(f"No description for frame {frame_num_response}.")
365 elif frame_num_response < 0 or frame_num_response >= len(frame_nums):
366 self.logger.warning(
367 f"Invalid frame_num {frame_num_response} in response: {response_line_strip}.")
368 else:
369 frame_num = frame_nums[frame_num_response]
370 frame_description_path = f"{self.job_path}/frame_{frame_num:04d}.txt"
371 async with aiofiles.open(frame_description_path, "w") as f:
372 await f.write(description)
373 num_described_frames += 1
375 scene = self.find_scene_for_frame(frame_num)
376 if scene:
377 scene.add_description(description)
378 self.logger.info(
379 f" Frame {frame_num} in scene {scene.scene_id}: "
380 f"{description[0:MAX_LOG_TEXT]}...")
381 else:
382 self.logger.warning(f"No scene found for frame {frame_num}.")
383 except Exception as ex:
384 self.logger.error(f"Error parsing description line '{response_line_strip}': {ex}")
386 if num_described_frames < len(frame_nums):
387 self.logger.warning(
388 f"Only described {num_described_frames} out of {len(frame_nums)} frames in batch {batch_id}.")
390 async def choose_scenes_for_highlight(
391 self,
392 total_length: int = 30,
393 max_scenes: int = 5,
394 max_tokens: int = 256,
395 ) -> List[int]:
396 """Select the scenes for highlight short."""
397 prompt = HIGHLIGHT_PROMPT.format(
398 total_length=total_length,
399 max_scenes=max_scenes,
400 )
401 for scene in self.scenes:
402 prompt += f"- Scene {scene.scene_id}:\n"
403 prompt += f"Duration: {scene.duration_sec:.1f} seconds.\n"
404 if scene.descriptions:
405 prompt += f"Description: {' '.join(scene.descriptions)}\n"
406 if scene.transcript:
407 prompt += f"Transcript: {scene.transcript}\n"
408 prompt += "\n"
410 messages = [
411 {"role": "user", "content": prompt}
412 ]
414 prompt_json_path = f"{self.job_path}/highlight_prompt.json"
415 async with aiofiles.open(prompt_json_path, "w") as prompt_file:
416 await prompt_file.write(json.dumps(messages, indent=2))
418 prompt_path = f"{self.job_path}/highlight_prompt.txt"
419 async with aiofiles.open(prompt_path, "w") as prompt_file:
420 messages_json = json.dumps(messages, indent=2)
421 await prompt_file.write(messages_json)
423 """
424 response = await self.llm_client.chat.completions.create(
425 model=self.llm_model,
426 messages=messages,
427 max_tokens=max_tokens,
428 extra_body=None,
429 # stream=True,
430 )
431 response_message = response.choices[0].message
432 response_message_strip = response_message.content.strip()
433 """
434 response_message_strip = await self.gen.gen_text(
435 messages=messages,
436 max_tokens=max_tokens,
437 task_id="highlight",
438 )
440 response_json = []
441 try:
442 response_json = json.loads(response_message_strip)
443 except Exception as ex:
444 self.logger.error(f"Error parsing {response_message_strip}: {ex}")
445 return response_json
447 async def save_selected_scenes(
448 self,
449 chosen_scenes: List[int],
450 ) -> None:
451 """Persist the chosen scene IDs to selected_scenes.json for the WebUI.
453 Args:
454 chosen_scenes: Zero-based indices into self.scenes that were
455 selected by choose_scenes_for_highlight().
456 """
457 selected_scenes_path = f"{self.job_path}/selected_scenes.json"
458 async with aiofiles.open(selected_scenes_path, "w") as sel_file:
459 await sel_file.write(json.dumps(chosen_scenes, indent=2))
461 async def save_highlight_short(
462 self,
463 chosen_scenes: List[int]
464 ) -> str:
465 """Save the highlight short video."""
466 input_video_path = f"{self.job_path}/video.mp4"
467 input_video_binary = await read_file_bytes(input_video_path)
469 highlight_scene_binaries = []
470 for chosen_scene_id in chosen_scenes:
471 scene = self.scenes[chosen_scene_id]
472 scene_binary = chunk_video_binary(
473 video_binary=input_video_binary,
474 start_seconds=scene.start_sec,
475 end_seconds=scene.end_sec,
476 # Resize to the requested size if needed
477 width=self.width,
478 height=self.height,
479 )
480 highlight_scene_binaries.append(scene_binary)
482 out_video_binary = await concatenate_videos(highlight_scene_binaries)
483 out_video_path = f"{self.job_path}/{self.job_id}.mp4"
484 async with aiofiles.open(out_video_path, "wb") as file:
485 await file.write(out_video_binary)
486 return out_video_path
488 async def chunk_audio_into_scenes(self) -> List[str]:
489 """
490 Chunk the audio of the video into scenes.
491 """
492 chunks = []
493 try:
494 video_path = f"{self.job_path}/video.mp4"
495 audio_path = f"{self.job_path}/audio.wav"
496 audio_path = await extract_audio_from_video(video_path, audio_path)
497 audio_base64 = await read_file_base64(audio_path)
498 self.logger.info(f"Extracted audio with {bytes_to_human(len(audio_base64))}.")
500 for scene in self.scenes:
501 scene_audio_base64 = chunk_audio_base64(
502 audio_base64=audio_base64,
503 start_seconds=scene.start_sec,
504 end_seconds=scene.end_sec)
505 scene_audio_path = f"{self.job_path}/scene_{scene.scene_id:03d}.wav"
506 await save_base64_as_binary(scene_audio_path, scene_audio_base64)
507 scene.audio_path = f"scene_{scene.scene_id:03d}.wav"
508 chunks.append(scene_audio_base64)
509 except Exception as ex:
510 self.logger.error(f"Error during audio chunking: {ex} [{type(ex)}]")
511 return chunks
513 async def transcribe_audio(
514 self
515 ) -> str:
516 """
517 Transcribe the audio of the video.
518 """
519 ret = ""
520 try:
521 for scene in self.scenes:
522 if not scene.audio_path:
523 continue
524 audio_path = f"{self.job_path}/{scene.audio_path}"
525 audio_transcript, language_code = await self.gen.gen_audio_transcript(
526 audio_path,
527 task_id=f"{scene.scene_id:03d}",
528 )
529 if not audio_transcript:
530 continue
531 ret += audio_transcript + "\n"
532 scene.transcript = audio_transcript
533 scene.language = language_code
534 self.logger.info(f"Scene {scene.scene_id} transcript: {audio_transcript[0:80]}...")
535 transcript_path = f"{self.job_path}/scene_{scene.scene_id:03d}.txt"
536 async with aiofiles.open(transcript_path, "w") as file:
537 await file.write(audio_transcript)
538 except Exception as ex:
539 self.logger.error(f"Error during transcription: {ex} [{type(ex)}]")
540 return ret