Update Fight_detec_func.py

Fight_detec_func.py

@@ -1,103 +1,112 @@
 import tensorflow as tf
-from frame_slicer import extract_video_frames
-import cv2
 import os
 import numpy as np
+import cv2
+import logging
+from datetime import datetime
 import matplotlib.pyplot as plt
 
-# Configuration
-import os
-MODEL_PATH = os.path.join(os.path.dirname(__file__), "trainnig_output", "final_model_2.h5")
+from frame_slicer import extract_video_frames  # Make sure this module is in PYTHONPATH
+
+# ----------------- Configuration ----------------- #
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "training_output", "final_model_2.h5")
 N_FRAMES = 30
 IMG_SIZE = (96, 96)
-RESULT_PATH = os.path.join(os.path.dirname(__file__), "results")  # Will be created if doesn't exist
+RESULT_PATH = os.path.join(os.path.dirname(__file__), "results")
+FIGHT_THRESHOLD = 0.61
 
-def fight_detec(video_path: str, debug: bool = True):
-    """Detects fight in a video and returns the result and confidence score."""
-    
-    class FightDetector:
-        def __init__(self):
-            self.model = self._load_model()
-        
-        def _load_model(self):
-            try:
-                model = tf.keras.models.load_model(MODEL_PATH, compile=False)
-                if debug:
-                    print("\nModel loaded successfully. Input shape:", model.input_shape)
-                return model
-            except Exception as e:
-                print(f"Model loading failed: {e}")
-                return None
-        
-        def _extract_frames(self, video_path):
-            frames = extract_video_frames(video_path, N_FRAMES, IMG_SIZE)
+os.makedirs(RESULT_PATH, exist_ok=True)
+
+# ----------------- Logging Setup ----------------- #
+logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')
+logger = logging.getLogger(__name__)
+
+# ----------------- Main Detector Class ----------------- #
+class FightDetector:
+    def __init__(self):
+        self.model = self._load_model()
+
+    def _load_model(self):
+        try:
+            model = tf.keras.models.load_model(MODEL_PATH, compile=False)
+            logger.info(f"Model loaded successfully. Input shape: {model.input_shape}")
+            return model
+        except Exception as e:
+            logger.error(f"Model loading failed: {e}")
+            return None
+
+    def _extract_frames(self, video_path):
+        frames = extract_video_frames(video_path, N_FRAMES, IMG_SIZE)
+        if frames is None:
+            return None
+
+        blank_frames = np.all(frames == 0, axis=(1, 2, 3)).sum()
+        if blank_frames > 0:
+            logger.warning(f"{blank_frames} blank frames detected.")
+
+        sample_frame = (frames[0] * 255).astype(np.uint8)
+        debug_frame_path = os.path.join(RESULT_PATH, 'debug_frame.jpg')
+        cv2.imwrite(debug_frame_path, cv2.cvtColor(sample_frame, cv2.COLOR_RGB2BGR))
+        logger.info(f"Debug frame saved to: {debug_frame_path}")
+
+        return frames
+
+    def predict(self, video_path):
+        if not os.path.exists(video_path):
+            return "Error: Video not found", None
+
+        try:
+            frames = self._extract_frames(video_path)
             if frames is None:
-                return None
-            
-            if debug:
-                blank_frames = np.all(frames == 0, axis=(1, 2, 3)).sum()
-                if blank_frames > 0:
-                    print(f"Warning: {blank_frames} blank frames detected")
-                sample_frame = (frames[0] * 255).astype(np.uint8)
-                os.makedirs(RESULT_PATH, exist_ok=True)
-                cv2.imwrite(os.path.join(RESULT_PATH, 'debug_frame.jpg'), 
-                            cv2.cvtColor(sample_frame, cv2.COLOR_RGB2BGR))
-            
-            return frames
-        
-        def predict(self, video_path):
-            if not os.path.exists(video_path):
-                return "Error: Video not found", None
-            
-            try:
-                frames = self._extract_frames(video_path)
-                if frames is None:
-                    return "Error: Frame extraction failed", None
-                
-                if frames.shape[0] != N_FRAMES:
-                    return f"Error: Expected {N_FRAMES} frames, got {frames.shape[0]}", None
-                
-                if np.all(frames == 0):
-                    return "Error: All frames are blank", None
-                
-                prediction = self.model.predict(frames[np.newaxis, ...], verbose=0)[0][0]
-                result = "FIGHT" if prediction >= 0.61 else "NORMAL"
-                confidence = min(max(abs(prediction - 0.61) * 150 + 50, 0), 100)
-                
-                if debug:
-                    self._debug_visualization(frames, prediction, result, video_path)
-                
-                return f"{result} ({confidence:.1f}% confidence)", prediction
-            
-            except Exception as e:
-                return f"Prediction error: {str(e)}", None
-
-        def _debug_visualization(self, frames, score, result, video_path):
-            print(f"\nPrediction Score: {score:.4f}")
-            print(f"Decision: {result}")
-            plt.figure(figsize=(15, 5))
-            for i in range(min(10, len(frames))):
-                plt.subplot(2, 5, i+1)
-                plt.imshow(frames[i])
-                plt.title(f"Frame {i}\nMean: {frames[i].mean():.2f}")
-                plt.axis('off')
-            plt.suptitle(f"Prediction: {result} (Score: {score:.4f})")
-            plt.tight_layout()
-
-            # Save the visualization
-            base_name = os.path.splitext(os.path.basename(video_path))[0]
-            save_path = os.path.join(RESULT_PATH, f"{base_name}_prediction_result.png")
-            plt.savefig(save_path)
-            plt.close()
-            print(f"Visualization saved to: {save_path}")
+                return "Error: Frame extraction failed", None
+
+            if frames.shape[0] != N_FRAMES:
+                return f"Error: Expected {N_FRAMES} frames, got {frames.shape[0]}", None
+
+            if np.all(frames == 0):
+                return "Error: All frames are blank", None
+
+            prediction = self.model.predict(frames[np.newaxis, ...], verbose=0)[0][0]
+            result = "FIGHT" if prediction >= FIGHT_THRESHOLD else "NORMAL"
+            confidence = min(max(abs(prediction - FIGHT_THRESHOLD) * 150 + 50, 0), 100)
+
+            self._debug_visualization(frames, prediction, result, video_path)
 
+            return f"{result} ({confidence:.1f}% confidence)", prediction
+
+        except Exception as e:
+            return f"Prediction error: {str(e)}", None
+
+    def _debug_visualization(self, frames, score, result, video_path):
+        logger.info(f"Prediction Score: {score:.4f}")
+        logger.info(f"Decision: {result}")
+        plt.figure(figsize=(15, 5))
+        for i in range(min(10, len(frames))):
+            plt.subplot(2, 5, i + 1)
+            plt.imshow(frames[i])
+            plt.title(f"Frame {i}\nMean: {frames[i].mean():.2f}")
+            plt.axis('off')
+        plt.suptitle(f"Prediction: {result} (Score: {score:.4f})")
+        plt.tight_layout()
+
+        base_name = os.path.splitext(os.path.basename(video_path))[0]
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        save_path = os.path.join(RESULT_PATH, f"{base_name}_result_{timestamp}.png")
+        plt.savefig(save_path)
+        plt.close()
+        logger.info(f"Visualization saved to: {save_path}")
+
+# ----------------- External Interface ----------------- #
+def fight_detec(video_path: str, debug: bool = True):
     detector = FightDetector()
     if detector.model is None:
         return "Error: Model loading failed", None
     return detector.predict(video_path)
 
-# # Entry point
-# path0 = input("Enter the local path to the video file to detect fight: ")
-# path = path0.strip('"')  # Remove extra quotes if copied from Windows
-# print(f"[INFO] Loading video: {path}")
-# fight_detec(path)
+# ----------------- Optional CLI Entry ----------------- #
+if __name__ == "__main__":
+    path0 = input("Enter the local path to the video file to detect fight: ")
+    path = path0.strip('"')  # Remove extra quotes if copied from Windows
+    logger.info(f"Loading video: {path}")
+    result, score = fight_detec(path)
+    logger.info(f"Result: {result}")