Update Fight_detec_func.py

Fight_detec_func.py

@@ -1,112 +1,189 @@
+--- START OF FILE Fight_detec_func.py ---
+
 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
 
-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")
+# Configuration
+import os
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "trainnig_output", "final_model_2.h5")
 N_FRAMES = 30
 IMG_SIZE = (96, 96)
+# Define RESULT_PATH relative to the script location
 RESULT_PATH = os.path.join(os.path.dirname(__file__), "results")
-FIGHT_THRESHOLD = 0.61
-
-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 "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):
+    """Detects fight in a video and returns the result string and raw prediction score."""
+
+    class FightDetector:
+        def __init__(self):
+            self.model = self._load_model()
+
+        def _load_model(self):
+            # Ensure the model path exists before loading
+            if not os.path.exists(MODEL_PATH):
+                print(f"Error: Model file not found at {MODEL_PATH}")
+                return None
+            try:
+                # Load model with compile=False if optimizer state isn't needed for inference
+                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)
+            if frames is None:
+                print(f"Frame extraction returned None for {video_path}")
+                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")
+                # Save a sample frame for debugging only if debug is True
+                if frames.shape[0] > 0 and not np.all(frames[0] == 0): # Avoid saving blank frame
+                    sample_frame = (frames[0] * 255).astype(np.uint8)
+                    try:
+                        os.makedirs(RESULT_PATH, exist_ok=True) # Ensure result path exists
+                        debug_frame_path = os.path.join(RESULT_PATH, 'debug_frame.jpg')
+                        cv2.imwrite(debug_frame_path, cv2.cvtColor(sample_frame, cv2.COLOR_RGB2BGR))
+                        print(f"Debug frame saved to {debug_frame_path}")
+                    except Exception as e:
+                        print(f"Failed to save debug frame: {e}")
+                else:
+                     print("Skipping debug frame save (first frame blank or no frames).")
+
+
+            return frames
+
+        def predict(self, video_path):
+            if not os.path.exists(video_path):
+                print(f"Error: Video not found at {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:
+                    # Pad with last frame or zeros if not enough frames were extracted
+                    print(f"Warning: Expected {N_FRAMES} frames, got {frames.shape[0]}. Padding...")
+                    if frames.shape[0] == 0: # No frames at all
+                         frames = np.zeros((N_FRAMES, *IMG_SIZE, 3), dtype=np.float32)
+                    else: # Pad with the last available frame
+                        padding_needed = N_FRAMES - frames.shape[0]
+                        last_frame = frames[-1][np.newaxis, ...]
+                        padding = np.repeat(last_frame, padding_needed, axis=0)
+                        frames = np.concatenate((frames, padding), axis=0)
+                    print(f"Frames padded to shape: {frames.shape}")
+
+
+                if np.all(frames == 0):
+                    # Check if all frames are actually blank (can happen with padding)
+                    print("Error: All frames are blank after processing/padding.")
+                    return "Error: All frames are blank", None
+
+                # Perform prediction
+                prediction = self.model.predict(frames[np.newaxis, ...], verbose=0)[0][0]
+                # Determine result based on threshold
+                threshold = 0.61 # Example threshold
+                is_fight = prediction >= threshold
+                result = "FIGHT" if is_fight else "NORMAL"
+
+                # Calculate confidence (simple distance from threshold, scaled)
+                # Adjust scaling factor (e.g., 150) and base (e.g., 50) as needed
+                # Ensure confidence reflects certainty (higher for values far from threshold)
+                if is_fight:
+                    confidence = min(max((prediction - threshold) * 150 + 50, 0), 100)
+                else:
+                    confidence = min(max((threshold - prediction) * 150 + 50, 0), 100)
+
+
+                result_string = f"{result} ({confidence:.1f}% confidence)"
+
+                if debug:
+                    print(f"Raw Prediction Score: {prediction:.4f}")
+                    self._debug_visualization(frames, prediction, result_string, video_path)
+
+                return result_string, float(prediction) # Return string and raw score
+
+            except Exception as e:
+                print(f"Prediction error: {str(e)}")
+                # Consider logging the full traceback here in a real application
+                # import traceback
+                # print(traceback.format_exc())
+                return f"Prediction error: {str(e)}", None
+
+        def _debug_visualization(self, frames, score, result, video_path):
+            # This function will only run if debug=True is passed to fight_detec
+            print(f"\n--- Debug Visualization ---")
+            print(f"Prediction Score: {score:.4f}")
+            print(f"Decision: {result}")
+
+            # Avoid plotting if matplotlib is not available or causes issues in deployment
+            try:
+                import matplotlib.pyplot as plt
+                plt.figure(figsize=(15, 5))
+                num_frames_to_show = min(10, len(frames))
+                for i in range(num_frames_to_show):
+                    plt.subplot(2, 5, i+1)
+                    # Ensure frame values are valid for imshow (0-1 or 0-255)
+                    img_display = frames[i]
+                    if np.max(img_display) <= 1.0: # Assuming normalized float [0,1]
+                         img_display = (img_display * 255).astype(np.uint8)
+                    else: # Assuming it might already be uint8 [0,255]
+                         img_display = img_display.astype(np.uint8)
+
+                    plt.imshow(img_display)
+                    plt.title(f"Frame {i}\nMean: {frames[i].mean():.2f}") # Use original frame for mean
+                    plt.axis('off')
+                plt.suptitle(f"Video: {os.path.basename(video_path)}\nPrediction: {result} (Raw Score: {score:.4f})")
+                plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout
+
+                # Save the visualization
+                os.makedirs(RESULT_PATH, exist_ok=True) # Ensure result path exists again
+                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() # Close the plot to free memory
+                print(f"Debug visualization saved to: {save_path}")
+            except ImportError:
+                print("Matplotlib not found. Skipping debug visualization plot.")
+            except Exception as e:
+                print(f"Error during debug visualization: {e}")
+            print("--- End Debug Visualization ---")
+
+
+    # --- Main function logic ---
     detector = FightDetector()
     if detector.model is None:
+        # Model loading failed, return error
         return "Error: Model loading failed", None
-    return detector.predict(video_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}")
+
+    # Call the predict method
+    result_str, prediction_score = detector.predict(video_path)
+    return result_str, prediction_score
+
+
+# # Example usage (commented out for library use)
+# if __name__ == "__main__":
+#     # Example of how to call the function
+#     test_video = input("Enter the local path to the video file: ").strip('"')
+#     if os.path.exists(test_video):
+#         print(f"[INFO] Processing video: {test_video}")
+#         result, score = fight_detec(test_video, debug=True) # Enable debug for local testing
+#         print(f"\nFinal Result: {result}")
+#         if score is not None:
+#             print(f"Raw Score: {score:.4f}")
+#     else:
+#         print(f"Error: File not found - {test_video}")
+
+--- END OF FILE Fight_detec_func.py ---