jhj0517 commited on
Commit
87304d5
·
1 Parent(s): 81226a3

Add cleaning logic

Browse files
Files changed (1) hide show
  1. modules/utils/video_helper.py +8 -2
modules/utils/video_helper.py CHANGED
@@ -23,11 +23,15 @@ class VideoInfo:
23
  def extract_frames(
24
  vid_input: str,
25
  output_temp_dir: str = TEMP_VIDEO_FRAMES_DIR,
26
- start_number: int = 0
 
27
  ):
28
  """
29
  Extract frames as jpg files and save them into output_temp_dir. This needs FFmpeg installed.
30
  """
 
 
 
31
  os.makedirs(output_temp_dir, exist_ok=True)
32
  output_path = os.path.join(output_temp_dir, "%05d.jpg")
33
 
@@ -258,7 +262,9 @@ def clean_temp_dir(temp_dir: Optional[str] = None):
258
 
259
  clean_files_with_extension(temp_dir, SOUND_FILE_EXT)
260
  clean_files_with_extension(temp_dir, IMAGE_FILE_EXT)
261
- clean_files_with_extension(temp_out_dir, IMAGE_FILE_EXT)
 
 
262
 
263
 
264
  def clean_files_with_extension(dir_path: str, extensions: List):
 
23
  def extract_frames(
24
  vid_input: str,
25
  output_temp_dir: str = TEMP_VIDEO_FRAMES_DIR,
26
+ start_number: int = 0,
27
+ clean=True
28
  ):
29
  """
30
  Extract frames as jpg files and save them into output_temp_dir. This needs FFmpeg installed.
31
  """
32
+ if clean:
33
+ clean_temp_dir(temp_dir=output_temp_dir)
34
+
35
  os.makedirs(output_temp_dir, exist_ok=True)
36
  output_path = os.path.join(output_temp_dir, "%05d.jpg")
37
 
 
262
 
263
  clean_files_with_extension(temp_dir, SOUND_FILE_EXT)
264
  clean_files_with_extension(temp_dir, IMAGE_FILE_EXT)
265
+
266
+ if os.path.exists(temp_out_dir):
267
+ clean_files_with_extension(temp_out_dir, IMAGE_FILE_EXT)
268
 
269
 
270
  def clean_files_with_extension(dir_path: str, extensions: List):