Refactor to `clean_files_with_extension()`

modules/sam_inference.py

@@ -15,7 +15,7 @@ from modules.model_downloader import (
     download_sam_model_url
 )
 from modules.paths import SAM2_CONFIGS_DIR, MODELS_DIR, TEMP_OUT_DIR, TEMP_DIR
-from modules.constants import BOX_PROMPT_MODE, AUTOMATIC_MODE, COLOR_FILTER, PIXELIZE_FILTER
+from modules.constants import BOX_PROMPT_MODE, AUTOMATIC_MODE, COLOR_FILTER, PIXELIZE_FILTER, IMAGE_FILE_EXT
 from modules.mask_utils import (
     save_psd_with_masks,
     create_mask_combined_images,
@@ -24,7 +24,7 @@ from modules.mask_utils import (
     create_solid_color_mask_image
 )
 from modules.video_utils import (get_frames_from_dir, create_video_from_frames, get_video_info, extract_frames,
-                                 extract_sound, clean_temp_dir, clean_image_files)
+                                 extract_sound, clean_temp_dir, clean_files_with_extension)
 from modules.utils import save_image
 from modules.logger_util import get_logger
 
@@ -277,14 +277,14 @@ class SamInference:
             logger.error(error_message)
             raise gr.Error(error_message, duration=20)
 
-        clean_image_files(TEMP_OUT_DIR)
+        clean_files_with_extension(TEMP_OUT_DIR, IMAGE_FILE_EXT)
+        self.video_predictor.reset_state(self.video_inference_state)
 
         prompt_frame_image, prompt = image_prompt_input_data["image"], image_prompt_input_data["points"]
 
         point_labels, point_coords, box = self.handle_prompt_data(prompt)
         obj_id = frame_idx
 
-        self.video_predictor.reset_state(self.video_inference_state)
         idx, scores, logits = self.add_prediction_to_frame(
             frame_idx=frame_idx,
             obj_id=obj_id,

modules/video_utils.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 import re
 
 from modules.logger_util import get_logger
+from modules.constants import SOUND_FILE_EXT, VIDEO_FILE_EXT, IMAGE_FILE_EXT
 from modules.paths import TEMP_DIR, TEMP_OUT_DIR
 
 logger = get_logger()
@@ -222,24 +223,12 @@ def clean_temp_dir(temp_dir: Optional[str] = None):
     else:
         temp_out_dir = os.path.join(temp_dir, "out")
 
-    clean_sound_files(temp_dir)
-    clean_image_files(temp_dir)
-    clean_image_files(temp_out_dir)
+    clean_files_with_extension(temp_dir, SOUND_FILE_EXT)
+    clean_files_with_extension(temp_dir, IMAGE_FILE_EXT)
+    clean_files_with_extension(temp_out_dir, IMAGE_FILE_EXT)
 
 
-def clean_sound_files(sound_dir: str):
-    """Removes all sound files from the directory."""
-    sound_extensions = ['.mp3', '.wav', '.aac', '.flac', '.ogg', '.m4a', '.wma']
-    _clean_files_with_extension(sound_dir, sound_extensions)
-
-
-def clean_image_files(image_dir: str):
-    """Removes all image files from the dir"""
-    image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp']
-    _clean_files_with_extension(image_dir, image_extensions)
-
-
-def _clean_files_with_extension(dir_path: str, extensions: List):
+def clean_files_with_extension(dir_path: str, extensions: List):
     for filename in os.listdir(dir_path):
         if filename.lower().endswith(tuple(extensions)):
             file_path = os.path.join(dir_path, filename)