diff --git a/app/main.py b/app/main.py index d9fc3f8..9e3e8fa 100644 --- a/app/main.py +++ b/app/main.py @@ -116,55 +116,49 @@ def process_cloud_storage_video(self): # Proceed to run the extraction process self.run_extraction(video_path, unique_filename) - # def run_extraction(self, video_path, unique_filename): - # class_config_path = os.path.join(self.config.object_class_directory, self.class_config_selection) - # specific_output_dir = os.path.join(self.config.output_directory, unique_filename) - # os.makedirs(specific_output_dir, exist_ok=True) - # output_format_instance = self.format_options[self.format_selection](output_dir=specific_output_dir, - # sahi_enabled=self.sahi_enabled) - # try: - # extractor = VideoFrameExtractor(self.config, video_path, self.frame_rate, specific_output_dir, - # self.model_selection, class_config_path, output_format_instance, - # self.transformations, self.model_types, self.sahi_config) - # extractor.extract_frames(self.model_confidence) - # if self.format_selection == "CVAT": - # output_format_instance.zip_and_cleanup() - # if self.storage_option == 'Object Storage': - # self.upload_outputs(specific_output_dir) - # - # # Clean up: Remove the temporary video file after processing - # if os.path.exists(video_path): - # os.remove(video_path) - # print(f"Deleted temporary video file: {video_path}") - # - # st.success('Extraction Completed!') - # except Exception as e: - # st.error(f"An error occurred during frame extraction: {str(e)}") - def run_extraction(self, video_path, unique_filename): + """Handles the frame extraction process with conditional error handling based on debug mode.""" + + # Prepare paths and configurations class_config_path = os.path.join(self.config.object_class_directory, self.class_config_selection) specific_output_dir = os.path.join(self.config.output_directory, unique_filename) os.makedirs(specific_output_dir, exist_ok=True) - output_format_instance = self.format_options[self.format_selection](output_dir=specific_output_dir, - sahi_enabled=self.sahi_enabled) + output_format_instance = self.format_options[self.format_selection]( + output_dir=specific_output_dir, sahi_enabled=self.sahi_enabled) - extractor = VideoFrameExtractor(self.config, video_path, self.frame_rate, specific_output_dir, - self.model_selection, class_config_path, output_format_instance, - self.transformations, self.model_types, self.sahi_config) - extractor.extract_frames(self.model_confidence) + def extraction_logic(): + """Core logic for video frame extraction and post-processing.""" + extractor = VideoFrameExtractor( + self.config, video_path, self.frame_rate, specific_output_dir, + self.model_selection, class_config_path, output_format_instance, + self.transformations, self.model_types, self.sahi_config) - if self.format_selection == "CVAT": - output_format_instance.zip_and_cleanup() + extractor.extract_frames(self.model_confidence) - if self.storage_option == 'Object Storage': - self.upload_outputs(specific_output_dir) + # Format-specific post-processing (e.g., zipping for CVAT format) + if self.format_selection == "CVAT": + output_format_instance.zip_and_cleanup() - # Clean up: Remove the temporary video file after processing - if os.path.exists(video_path): - os.remove(video_path) - print(f"Deleted temporary video file: {video_path}") + # Upload to object storage if configured + if self.storage_option == 'Object Storage': + self.upload_outputs(specific_output_dir) - st.success('Extraction Completed!') + # Clean up the temporary video file after processing + if os.path.exists(video_path): + os.remove(video_path) + print(f"Deleted temporary video file: {video_path}") + + # Notify user of successful extraction + st.success('Extraction Completed!') + + # Conditionally apply try-except block based on debug mode + if self.config.debug: + extraction_logic() # No error handling in debug mode + else: + try: + extraction_logic() # Error handling in production mode + except Exception as e: + st.error(f"An error occurred during frame extraction: {str(e)}") def upload_outputs(self, directory): """