latest changes

app.py

@@ -1,206 +1,187 @@
 import gradio as gr
 import torch
-from transformers import AutoImageProcessor, TimesformerForVideoClassification
 import cv2
-from PIL import Image
 import numpy as np
-import time
-from collections import deque
-import base64
-import io
-
-# --- Configuration ---
-# CHANGED: Using a public Facebook TimesFormer model fine-tuned on Kinetics
-HF_MODEL_REPO_ID = "facebook/timesformer-base-finetuned-k400"
-
-MODEL_INPUT_NUM_FRAMES = 8
-TARGET_IMAGE_HEIGHT = 224
-TARGET_IMAGE_WIDTH = 224
-RAW_RECORDING_DURATION_SECONDS = 10.0
-FRAMES_TO_SAMPLE_PER_CLIP = 20
-DELAY_BETWEEN_PREDICTIONS_SECONDS = 120.0 # 2 minutes for CPU, adjust for GPU
-
-# --- Load Model and Processor ---
-print(f"Loading model and processor from {HF_MODEL_REPO_ID}...")
-try:
-    processor = AutoImageProcessor.from_pretrained(HF_MODEL_REPO_ID)
-    model = TimesformerForVideoClassification.from_pretrained(HF_MODEL_REPO_ID)
-except Exception as e:
-    print(f"Error loading model: {e}")
-    exit()
-
-model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+import os
+import json
+from PIL import Image
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+import tempfile # For temporary file handling
+
+# --- 1. Define Model Architecture (Copy from small_video_classifier.py) ---
+# This is crucial because we need the model class definition to load weights.
+class SmallVideoClassifier(torch.nn.Module):
+    def __init__(self, num_classes=2, num_frames=8):
+        super(SmallVideoClassifier, self).__init__()
+        from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights
+        # Load weights only if they are available (they should be for IMAGENET1K_V1)
+        # Use a check to prevent error if weights are not found in specific environments
+        try:
+            weights = MobileNet_V3_Small_Weights.IMAGENET1K_V1
+        except Exception:
+            print("Warning: MobileNet_V3_Small_Weights.IMAGENET1K_V1 not found, initializing without pre-trained weights.")
+            weights = None # Or provide specific default for your use case
+
+        self.feature_extractor = mobilenet_v3_small(weights=weights)
+        self.feature_extractor.classifier = torch.nn.Identity()
+        self.num_spatial_features = 576 # MobileNetV3-Small's final feature map size
+        self.temporal_aggregator = torch.nn.AdaptiveAvgPool1d(1)
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Linear(self.num_spatial_features, 512),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(512, num_classes)
+        )
+
+    def forward(self, pixel_values):
+        batch_size, num_frames, channels, height, width = pixel_values.shape
+        x = pixel_values.view(batch_size * num_frames, channels, height, width)
+        spatial_features = self.feature_extractor(x)
+        spatial_features = spatial_features.view(batch_size, num_frames, self.num_spatial_features)
+        temporal_features = self.temporal_aggregator(spatial_features.permute(0, 2, 1)).squeeze(-1)
+        logits = self.classifier(temporal_features)
+        return logits
+
+# --- 2. Configuration and Model Loading ---
+HF_USERNAME = "owinymarvin"
+NEW_MODEL_REPO_ID_SHORT = "timesformer-violence-detector"
+NEW_MODEL_REPO_ID = f"{HF_USERNAME}/{NEW_MODEL_REPO_ID_SHORT}"
+
+# Download config.json to get model parameters
+print(f"Downloading config.json from {NEW_MODEL_REPO_ID}...")
+config_path = hf_hub_download(repo_id=NEW_MODEL_REPO_ID, filename="config.json")
+with open(config_path, 'r') as f:
+    model_config = json.load(f)
+
+NUM_FRAMES = model_config.get('num_frames', 8)
+IMAGE_SIZE = tuple(model_config.get('image_size', [224, 224]))
+NUM_CLASSES = model_config.get('num_classes', 2)
+
+# Define class labels (adjust if your dataset had different labels/order)
+CLASS_LABELS = ["Non-violence", "Violence"]
+if NUM_CLASSES != len(CLASS_LABELS):
+    print(f"Warning: NUM_CLASSES in config ({NUM_CLASSES}) does not match hardcoded CLASS_LABELS length ({len(CLASS_LABELS)}). Adjust CLASS_LABELS if needed.")
+
+
+# Initialize the model
+device = torch.device("cpu") # Explicitly use CPU as requested
+print(f"Using device: {device}")
+
+model = SmallVideoClassifier(num_classes=NUM_CLASSES, num_frames=NUM_FRAMES)
+
+# Download model weights
+print(f"Downloading model weights from {NEW_MODEL_REPO_ID}...")
+model_weights_path = hf_hub_download(repo_id=NEW_MODEL_REPO_ID, filename="small_violence_classifier.pth")
+model.load_state_dict(torch.load(model_weights_path, map_location=device))
 model.to(device)
-print(f"Model loaded on {device}.")
-print(f"Model's class labels (Kinetics): {model.config.id2label}") # Print new labels
-
-# --- Global State Variables for Live Demo ---
-raw_frames_buffer = deque()
-current_clip_start_time = time.time()
-last_prediction_completion_time = time.time()
-app_state = "recording" # States: "recording", "predicting", "processing_delay"
-
-# --- Helper function to sample frames ---
-def sample_frames(frames_list, target_count):
-    if not frames_list:
-        return []
-    if len(frames_list) <= target_count:
-        return frames_list
-    indices = np.linspace(0, len(frames_list) - 1, target_count, dtype=int)
-    sampled = [frames_list[int(i)] for i in indices]
-    return sampled
-
-# --- Main processing function for Live Demo Stream ---
-def live_predict_stream(image_np_array):
-    global raw_frames_buffer, current_clip_start_time, last_prediction_completion_time, app_state
-
-    current_time = time.time()
-    pil_image = Image.fromarray(image_np_array)
-
-    if app_state == "recording":
-        raw_frames_buffer.append(pil_image)
-        elapsed_recording_time = current_time - current_clip_start_time
+model.eval() # Set model to evaluation mode
+
+print(f"Model loaded successfully with {NUM_FRAMES} frames and image size {IMAGE_SIZE}.")
+
+# --- 3. Define Preprocessing Transform ---
+transform = transforms.Compose([
+    transforms.Resize(IMAGE_SIZE),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# --- 4. Gradio Inference Function ---
+def predict_video(video_path):
+    if video_path is None:
+        return None # Or raise an error, or return a placeholder video
+
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        print(f"Error: Could not open video file {video_path}.")
+        raise ValueError(f"Could not open video file {video_path}. Please ensure it's a valid video format.")
+
+    # Get video properties
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    # Create a temporary output video file
+    # Use tempfile to ensure proper cleanup on Hugging Face Spaces
+    temp_output_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    output_video_path = temp_output_file.name
+    temp_output_file.close() # Close the file handle as cv2.VideoWriter needs to open it
+
+    # Define the codec and create VideoWriter object
+    # For MP4, 'mp4v' is generally compatible.
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+
+    print(f"Processing video: {video_path}")
+    print(f"Total frames: {total_frames}, FPS: {fps}")
+    print(f"Output video will be saved to: {output_video_path}")
+
+    frame_buffer = [] # To store NUM_FRAMES for each prediction batch
+    current_prediction_label = "Processing..." # Initial label
+
+    frame_idx = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break # End of video
+
+        frame_idx += 1
         
-        yield f"Recording: {elapsed_recording_time:.1f}/{RAW_RECORDING_DURATION_SECONDS}s. Raw frames: {len(raw_frames_buffer)}", "Buffering..."
-
-        if elapsed_recording_time >= RAW_RECORDING_DURATION_SECONDS:
-            # Transition to predicting state
-            app_state = "predicting"
-            yield "Preparing to predict...", "Processing..."
-            print("DEBUG: Transitioning to 'predicting' state.")
-
-    elif app_state == "predicting":
-        # Ensure this prediction block only runs once per cycle
-        if raw_frames_buffer: # Only proceed if there are frames to process
-            print("DEBUG: Starting prediction.")
-            try:
-                sampled_raw_frames = sample_frames(list(raw_frames_buffer), FRAMES_TO_SAMPLE_PER_CLIP)
-                frames_for_model = sample_frames(sampled_raw_frames, MODEL_INPUT_NUM_FRAMES)
-
-                if len(frames_for_model) < MODEL_INPUT_NUM_FRAMES:
-                    yield "Error during frame sampling.", f"Error: Not enough frames ({len(frames_for_model)}/{MODEL_INPUT_NUM_FRAMES}). Resetting."
-                    print(f"ERROR: Insufficient frames for model input: {len(frames_for_model)}/{MODEL_INPUT_NUM_FRAMES}. Resetting state.")
-                    app_state = "recording" # Reset state to start a new recording
-                    raw_frames_buffer.clear()
-                    current_clip_start_time = time.time()
-                    last_prediction_completion_time = time.time()
-                    return # Exit this stream call to wait for next frame or reset
-
-                processed_input = processor(images=frames_for_model, return_tensors="pt")
-                pixel_values = processed_input.pixel_values.to(device)
-
-                with torch.no_grad():
-                    outputs = model(pixel_values)
-                    logits = outputs.logits
-
-                predicted_class_id = logits.argmax(-1).item()
-                predicted_label = model.config.id2label.get(predicted_class_id, "Unknown")
-                confidence = torch.nn.functional.softmax(logits, dim=-1)[0][predicted_class_id].item()
-
-                prediction_result = f"Predicted: {predicted_label} (Confidence: {confidence:.2f})"
-                status_message = "Prediction complete."
-                print(f"DEBUG: Prediction Result: {prediction_result}")
-                
-                # Yield the prediction result immediately to ensure UI update
-                yield status_message, prediction_result
-
-                # Clear buffer and transition to delay AFTER yielding the prediction
-                raw_frames_buffer.clear()
-                last_prediction_completion_time = current_time
-                app_state = "processing_delay"
-                print("DEBUG: Transitioning to 'processing_delay' state.")
-
-            except Exception as e:
-                error_message = f"Error during prediction: {e}"
-                print(f"ERROR during prediction: {e}")
-                # Yield error to UI
-                yield "Prediction error.", error_message
-                app_state = "processing_delay" # Still go to delay state to prevent constant errors
-                raw_frames_buffer.clear() # Clear buffer to prevent re-processing same problematic frames
-
-    elif app_state == "processing_delay":
-        elapsed_delay = current_time - last_prediction_completion_time
+        # Convert frame from BGR (OpenCV) to RGB (PIL/PyTorch)
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(frame_rgb)
         
-        if elapsed_delay < DELAY_BETWEEN_PREDICTIONS_SECONDS:
-            # Continue yielding the delay message and the last prediction result
-            yield f"Delaying next prediction: {int(elapsed_delay)}/{int(DELAY_BETWEEN_PREDICTIONS_SECONDS)}s", gr.NO_VALUE
-        else:
-            # Delay is over, reset for new recording cycle
-            app_state = "recording"
-            current_clip_start_time = current_time
-            print("DEBUG: Transitioning back to 'recording' state.")
-            yield "Starting new recording...", "Ready for new prediction."
+        # Apply transformations and add to buffer
+        processed_frame = transform(pil_image) # shape: (C, H, W)
+        frame_buffer.append(processed_frame)
+
+        # Perform prediction when the buffer is full
+        if len(frame_buffer) == NUM_FRAMES:
+            # Stack the buffered frames and add a batch dimension
+            # Resulting shape: (1, NUM_FRAMES, C, H, W)
+            input_tensor = torch.stack(frame_buffer, dim=0).unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                outputs = model(input_tensor)
+                probabilities = torch.softmax(outputs, dim=1)
+                predicted_class_idx = torch.argmax(probabilities, dim=1).item()
+                current_prediction_label = f"Prediction: {CLASS_LABELS[predicted_class_idx]} (Prob: {probabilities[0, predicted_class_idx]:.2f})"
+            
+            # Reset buffer for the next non-overlapping window
+            frame_buffer = [] 
+            # Or, if you want sliding window (more continuous output but higher compute):
+            # frame_buffer = frame_buffer[1:] # e.g., slide by 1 frame
+
+        # Draw prediction text on the current frame
+        # The prediction will lag by NUM_FRAMES, as it's based on the previous batch.
+        # We display the last known prediction.
+        cv2.putText(frame, current_prediction_label, (10, 30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2, cv2.LINE_AA)
+
+        # Write the processed frame to the output video
+        out.write(frame)
+
+    # Release resources
+    cap.release()
+    out.release()
+    print(f"Video processing complete. Output saved to: {output_video_path}")
     
-    pass
-
-def reset_app_state_manual():
-    global raw_frames_buffer, current_clip_start_time, last_prediction_completion_time, app_state
-    raw_frames_buffer.clear()
-    current_clip_start_time = time.time()
-    last_prediction_completion_time = time.time()
-    app_state = "recording"
-    print("DEBUG: Manual reset triggered.")
-    # Return initial values immediately upon reset
-    return "Ready to record...", "Ready for new prediction."
-
-# --- Gradio UI Layout ---
-with gr.Blocks() as demo:
-    gr.Markdown(
-        f"""
-        # TimesFormer Action Recognition - Using Facebook Kinetics Model
-        This Space hosts the `{HF_MODEL_REPO_ID}` model.
-        Live webcam demo with recording and prediction phases.
-        **NOTE: This model predicts general human actions (e.g., 'playing guitar', 'walking'), not crime events.**
-        """
-    )
-
-    with gr.Tab("Live Webcam Demo"):
-        gr.Markdown(
-            f"""
-            Continuously captures live webcam feed for **{RAW_RECORDING_DURATION_SECONDS} seconds**,
-            then makes a prediction. There is a **{DELAY_BETWEEN_PREDICTIONS_SECONDS/60:.0f} minute delay** afterwards.
-            """
-        )
-        with gr.Row():
-            with gr.Column():
-                webcam_input = gr.Image(
-                    sources=["webcam"],
-                    streaming=True,
-                    label="Live Webcam Feed"
-                )
-                status_output = gr.Textbox(label="Current Status", value="Initializing...")
-                reset_button = gr.Button("Reset / Start New Cycle")
-            with gr.Column():
-                prediction_output = gr.Textbox(label="Prediction Result", value="Waiting...")
-
-        webcam_input.stream(
-            live_predict_stream,
-            inputs=[webcam_input],
-            outputs=[status_output, prediction_output]
-        )
-        
-        reset_button.click(
-            reset_app_state_manual,
-            inputs=[],
-            outputs=[status_output, prediction_output]
-        )
-
-    with gr.Tab("API Endpoint for External Clients"):
-        gr.Markdown(
-            """
-            Use this API endpoint to send base64-encoded frames for prediction.
-            (Currently uses the Kinetics model).
-            """
-        )
-        gr.Interface(
-            fn=lambda frames_list: "API endpoint is active for programmatic calls. See documentation in app.py.",
-            inputs=gr.Json(label="List of Base64-encoded image strings"),
-            outputs=gr.Textbox(label="API Response"),
-            live=False,
-            allow_flagging="never"
-        )
-
-
-if __name__ == "__main__":
-    demo.launch()
+    return output_video_path # Gradio expects the path to the output video
+
+# --- 5. Gradio Interface Setup ---
+iface = gr.Interface(
+    fn=predict_video,
+    inputs=gr.Video(label="Upload Video for Violence Detection (MP4 recommended)", type="filepath"),
+    outputs=gr.Video(label="Processed Video with Predictions"),
+    title="Real-time Violence Detection with SmallVideoClassifier",
+    description="Upload a video, and the model will analyze it for violence, displaying the predicted class and confidence on each frame.",
+    allow_flagging="never", # Disable flagging on Hugging Face Spaces
+    examples=[
+        # You can provide example video URLs or paths if you have them publicly available
+        # Example: "https://huggingface.co/datasets/gradio/test-files/resolve/main/video.mp4"
+    ]
+)
+
+iface.launch()

requirements.txt

@@ -1,5 +1,7 @@
 torch
-transformers
-opencv-python-headless # Use headless for server environments without display
+torchvision
+opencv-python-headless  # Use headless for server environments to avoid GUI dependencies
+gradio
+huggingface_hub
 Pillow
-gradio
+numpy