Update app.py

app.py

@@ -1,73 +1,150 @@
 import cv2
 import numpy as np
-from keras.models import load_model
-from moviepy.editor import VideoFileClip
+import tensorflow as tf
+from moviepy.editor import VideoFileClip, concatenate_videoclips
 import gradio as gr
+from tqdm import tqdm
+import os
 
-# Load the AI model
-model = load_model('model.h5')
+# --- Model Loading ---
+# Ensure you have a 'model.h5' file in the same directory.
+# This model should be trained to output a value close to 1 for a jumpscare frame
+# and close to 0 for a non-jumpscare frame.
+MODEL_PATH = 'model.h5'
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}. Please ensure a trained model exists.")
+model = tf.keras.models.load_model(MODEL_PATH)
 
-def process_frame(frame):
-    # Resize the image to (128, 128)
-    resized_frame = cv2.resize(frame, (128, 128))
-    resized_frame = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2RGB)
-    img_array = np.array(resized_frame) / 255.0
-    img_array = img_array.reshape(1, 128, 128, 3)
+
+def predict_frame_is_jumpscare(frame, threshold):
+    """
+    Analyzes a single video frame and predicts if it's a jumpscare.
     
-    prediction = model.predict(img_array)  # Get prediction
+    Args:
+        frame (np.array): A single video frame in BGR format from OpenCV.
+        threshold (float): The sensitivity threshold (0.0 to 1.0) for the prediction.
 
-    # Example logic to modify frame based on prediction
-    if prediction[0][0] > 0.5:  # Condition based on prediction
-        return frame  # Keep the original frame if the condition is met
-    else:
-        # Example modification: convert to grayscale
-        return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    Returns:
+        bool: True if the frame is predicted as a jumpscare, False otherwise.
+    """
+    # 1. Preprocess the frame for the model
+    # Convert BGR (OpenCV default) to RGB (model was likely trained on RGB)
+    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    # Resize to the model's expected input size (e.g., 128x128)
+    resized_frame = cv2.resize(rgb_frame, (128, 128))
+    # Normalize pixel values to be between 0 and 1
+    img_array = np.array(resized_frame) / 255.0
+    # Add a batch dimension (model expects shape: 1, height, width, channels)
+    img_array = np.expand_dims(img_array, axis=0)
 
-def convert_video_to_dumpscare(video, sensitivity):
-    output_frames = []
+    # 2. Make a prediction
+    prediction = model.predict(img_array, verbose=0) # verbose=0 silences a lot of Keras logs
     
-    # Read the video
-    cap = cv2.VideoCapture(video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    # 3. Return True/False based on the prediction and threshold
+    # We assume the model's output is a single value in prediction[0][0]
+    return prediction[0][0] > threshold
+
+
+def generate_jumpscare_compilation(video_path, sensitivity, progress=gr.Progress()):
+    """
+    Analyzes a video, finds jumpscare segments, and creates a compilation.
     
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter('output_dumpscare.mp4', fourcc, fps, (width, height))
+    Args:
+        video_path (str): The file path to the input video.
+        sensitivity (int): An integer from 0 to 100 for detection sensitivity.
+        progress (gr.Progress): Gradio progress tracker.
 
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+    Returns:
+        str or None: The file path to the output video, or None if no jumpscares are found.
+    """
+    # Convert sensitivity (0-100) to a model threshold (0.0-1.0)
+    threshold = sensitivity / 100.0
+    
+    # Use a lower FPS for analysis to speed things up. Jumpscares usually last
+    # for several frames, so we don't need to check every single one.
+    analysis_fps = 10 
+    
+    print("Loading video file...")
+    original_clip = VideoFileClip(video_path)
+    
+    jumpscare_segments = []
+    is_in_jumpscare_segment = False
+    segment_start_time = 0
+    
+    print(f"Analyzing video for jumpscares with threshold {threshold}...")
+    
+    # Use tqdm for a console progress bar and gr.Progress for the UI
+    total_frames = int(original_clip.duration * analysis_fps)
+    frame_iterator = original_clip.iter_frames(fps=analysis_fps)
+    
+    for i, frame in enumerate(tqdm(frame_iterator, total=total_frames, desc="Analyzing Frames")):
+        current_time = i / analysis_fps
+        progress(i / total_frames, desc=f"Analyzing... {int(current_time)}s / {int(original_clip.duration)}s")
+        
+        is_jumpscare = predict_frame_is_jumpscare(frame, threshold)
         
-        modified_frame = process_frame(frame)  # Get modified frame
-        out.write(modified_frame)  # Write the modified frame to the output video
+        # --- State machine to find continuous jumpscare segments ---
+        if is_jumpscare and not is_in_jumpscare_segment:
+            # Start of a new jumpscare segment
+            is_in_jumpscare_segment = True
+            segment_start_time = current_time
+        
+        elif not is_jumpscare and is_in_jumpscare_segment:
+            # End of the current jumpscare segment
+            is_in_jumpscare_segment = False
+            # Add a small buffer to the end time to capture the full scare
+            segment_end_time = current_time + 0.5 
+            jumpscare_segments.append((segment_start_time, segment_end_time))
+
+    # If the video ends while in a jumpscare, close the last segment
+    if is_in_jumpscare_segment:
+        jumpscare_segments.append((segment_start_time, original_clip.duration))
 
-    cap.release()
-    out.release()
+    # --- Create the final video ---
+    if not jumpscare_segments:
+        print("No jumpscares found.")
+        # We need to return a value that Gradio can handle.
+        # Returning None will cause an error. Instead, we can return a message.
+        # A better approach would be to update a gr.Text component, but for simplicity,
+        # we'll raise an error that Gradio can display.
+        raise gr.Error("No jumpscares were detected with the current sensitivity setting. Try a lower value.")
 
-    # Add audio from the original video
-    original_video = VideoFileClip(video)
-    output_video = VideoFileClip('output_dumpscare.mp4')
-    final_video = output_video.set_audio(original_video.audio)
-    final_video.write_videofile('final_output_dumpscare.mp4', codec='libx264')
+    print(f"Found {len(jumpscare_segments)} jumpscare segments. Creating compilation...")
+    progress(0.9, desc="Stitching clips together...")
 
-    return 'final_output_dumpscare.mp4'
+    # Create subclips from the original video using the detected timestamps
+    final_clips = [original_clip.subclip(start, end) for start, end in jumpscare_segments]
+    
+    # Concatenate all the jumpscare clips into one video
+    final_video = concatenate_videoclips(final_clips)
+    
+    output_path = "fnaf_jumpscare_compilation.mp4"
+    # Write the final video file, including audio
+    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac", temp_audiofile='temp-audio.m4a', remove_temp=True)
+    
+    # Close the clips to free up resources
+    original_clip.close()
+    for clip in final_clips:
+        clip.close()
+    
+    print(f"Compilation saved to {output_path}")
+    return output_path
 
-# Gradio Interface
-def gradio_interface(video, sensitivity):
-    output_video = convert_video_to_dumpscare(video, sensitivity)
-    return output_video
 
-# Create Gradio interface
+# --- Gradio Interface ---
+# Create a Gradio interface for the function
 iface = gr.Interface(
-    fn=gradio_interface,
+    fn=generate_jumpscare_compilation,
     inputs=[
-        gr.Video(label="Input Video"),
-        gr.Slider(label="Sensitivity", minimum=0, maximum=100, step=1)
+        gr.Video(label="Upload FNAF Video"),
+        gr.Slider(minimum=1, maximum=100, step=1, value=70, label="Detection Sensitivity", 
+                  info="Higher values mean the AI needs to be more certain it's a jumpscare. Lower values will detect more, but might have false positives.")
     ],
-    outputs=gr.Video(label="Output Video")
+    outputs=gr.Video(label="Jumpscare Compilation"),
+    title="AI FNAF Jumpscare Dump Generator",
+    description="Upload a video, and this tool will use an AI model to find all the jumpscares and compile them into a single video. This requires a `model.h5` file.",
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()