AC2513 commited on
Commit
67d411a
·
1 Parent(s): 45a53c4

adjusted option for max_images

Browse files
Files changed (1) hide show
  1. src/app.py +4 -6
src/app.py CHANGED
@@ -11,8 +11,6 @@ import cv2
11
  from loguru import logger
12
  from PIL import Image
13
 
14
- MAX_NUM_IMAGES = 6
15
-
16
  dotenv_path = find_dotenv()
17
 
18
  load_dotenv(dotenv_path)
@@ -28,7 +26,7 @@ model = Gemma3ForConditionalGeneration.from_pretrained(
28
  attn_implementation="eager",
29
  )
30
 
31
- def get_frames(video_path: str) -> list[tuple[Image.Image, float]]:
32
  capture = cv2.VideoCapture(video_path)
33
  if not capture.isOpened():
34
  raise ValueError(f"Could not open video file: {video_path}")
@@ -36,11 +34,11 @@ def get_frames(video_path: str) -> list[tuple[Image.Image, float]]:
36
  fps = capture.get(cv2.CAP_PROP_FPS)
37
  total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
38
 
39
- frame_interval = max(total_frames // MAX_NUM_IMAGES, 1)
40
  frames: list[tuple[Image.Image, float]] = []
41
 
42
- for i in range(0, min(total_frames, MAX_NUM_IMAGES * frame_interval), frame_interval):
43
- if len(frames) >= MAX_NUM_IMAGES:
44
  break
45
 
46
  capture.set(cv2.CAP_PROP_POS_FRAMES, i)
 
11
  from loguru import logger
12
  from PIL import Image
13
 
 
 
14
  dotenv_path = find_dotenv()
15
 
16
  load_dotenv(dotenv_path)
 
26
  attn_implementation="eager",
27
  )
28
 
29
+ def get_frames(video_path: str, max_images: int) -> list[tuple[Image.Image, float]]:
30
  capture = cv2.VideoCapture(video_path)
31
  if not capture.isOpened():
32
  raise ValueError(f"Could not open video file: {video_path}")
 
34
  fps = capture.get(cv2.CAP_PROP_FPS)
35
  total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
36
 
37
+ frame_interval = max(total_frames // max_images, 1)
38
  frames: list[tuple[Image.Image, float]] = []
39
 
40
+ for i in range(0, min(total_frames, max_images * frame_interval), frame_interval):
41
+ if len(frames) >= max_images:
42
  break
43
 
44
  capture.set(cv2.CAP_PROP_POS_FRAMES, i)