jhj0517 commited on
Commit
3ccfbcb
1 Parent(s): 4596138

Add codec to VideoInfo and update frames quality

Browse files
Files changed (1) hide show
  1. modules/video_utils.py +12 -5
modules/video_utils.py CHANGED
@@ -18,6 +18,7 @@ class VideoInfo:
18
  frame_rate: Optional[int] = None
19
  duration: Optional[float] = None
20
  has_sound: Optional[bool] = None
 
21
 
22
 
23
  def extract_frames(
@@ -35,6 +36,8 @@ def extract_frames(
35
  'ffmpeg',
36
  '-y', # Enable overwriting
37
  '-i', vid_input,
 
 
38
  '-start_number', str(start_number),
39
  f'{output_path}'
40
  ]
@@ -96,6 +99,7 @@ def get_video_info(vid_input: str) -> VideoInfo:
96
  frame_rate = None
97
  duration = None
98
  has_sound = False
 
99
 
100
  for line in output.splitlines():
101
  if 'Stream #0:0' in line and 'Video:' in line:
@@ -103,6 +107,10 @@ def get_video_info(vid_input: str) -> VideoInfo:
103
  if fps_match:
104
  frame_rate = float(fps_match.group(1))
105
 
 
 
 
 
106
  elif 'Duration:' in line:
107
  duration_match = re.search(r'Duration: (\d{2}):(\d{2}):(\d{2}\.\d{2})', line)
108
  if duration_match:
@@ -119,7 +127,8 @@ def get_video_info(vid_input: str) -> VideoInfo:
119
  num_frames=num_frames,
120
  frame_rate=frame_rate,
121
  duration=duration,
122
- has_sound=has_sound
 
123
  )
124
 
125
  except subprocess.CalledProcessError as e:
@@ -151,12 +160,12 @@ def create_video_from_frames(
151
  sound_path = temp_sound
152
 
153
  if frame_rate is None:
154
- frame_rate = 25 # Default frame rate
155
 
156
  command = [
157
  'ffmpeg',
158
  '-y',
159
- '-framerate', frame_rate,
160
  '-i', os.path.join(frames_dir, "%05d.jpg"),
161
  '-c:v', 'libx264',
162
  '-pix_fmt', 'yuv420p',
@@ -171,12 +180,10 @@ def create_video_from_frames(
171
  '-b:a', '192k',
172
  '-shortest'
173
  ]
174
- print(command)
175
  try:
176
  subprocess.run(command, check=True)
177
  except subprocess.CalledProcessError as e:
178
  logger.exception("Error occurred while creating video from frames")
179
- print("Done, output path:", output_path)
180
  return output_path
181
 
182
 
 
18
  frame_rate: Optional[int] = None
19
  duration: Optional[float] = None
20
  has_sound: Optional[bool] = None
21
+ codec: Optional[str] = None
22
 
23
 
24
  def extract_frames(
 
36
  'ffmpeg',
37
  '-y', # Enable overwriting
38
  '-i', vid_input,
39
+ '-qscale:v', '2',
40
+ '-vf', f'scale=iw:ih',
41
  '-start_number', str(start_number),
42
  f'{output_path}'
43
  ]
 
99
  frame_rate = None
100
  duration = None
101
  has_sound = False
102
+ codec = None
103
 
104
  for line in output.splitlines():
105
  if 'Stream #0:0' in line and 'Video:' in line:
 
107
  if fps_match:
108
  frame_rate = float(fps_match.group(1))
109
 
110
+ codec_match = re.search(r'Video: (\w+)', line)
111
+ if codec_match:
112
+ codec = codec_match.group(1)
113
+
114
  elif 'Duration:' in line:
115
  duration_match = re.search(r'Duration: (\d{2}):(\d{2}):(\d{2}\.\d{2})', line)
116
  if duration_match:
 
127
  num_frames=num_frames,
128
  frame_rate=frame_rate,
129
  duration=duration,
130
+ has_sound=has_sound,
131
+ codec=codec
132
  )
133
 
134
  except subprocess.CalledProcessError as e:
 
160
  sound_path = temp_sound
161
 
162
  if frame_rate is None:
163
+ frame_rate = 25 # Default frame rate for ffmpeg
164
 
165
  command = [
166
  'ffmpeg',
167
  '-y',
168
+ '-framerate', str(frame_rate),
169
  '-i', os.path.join(frames_dir, "%05d.jpg"),
170
  '-c:v', 'libx264',
171
  '-pix_fmt', 'yuv420p',
 
180
  '-b:a', '192k',
181
  '-shortest'
182
  ]
 
183
  try:
184
  subprocess.run(command, check=True)
185
  except subprocess.CalledProcessError as e:
186
  logger.exception("Error occurred while creating video from frames")
 
187
  return output_path
188
 
189