Merge pull request #9 from shamspias/feat/rt-detr

fix debug logic

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):
         """