Spaces:
Sleeping
Sleeping
| import tensorflow as tf | |
| from frame_slicer import extract_video_frames | |
| import cv2 | |
| import os | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| # 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) | |
| RESULT_PATH = os.path.join(os.path.dirname(__file__), "results") # Will be created if doesn't exist | |
| 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) | |
| 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}") | |
| 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) | |