jhj0517
commited on
Commit
·
87304d5
1
Parent(s):
81226a3
Add cleaning logic
Browse files
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 |
-
|
|
|
|
|
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):
|