app.py

import streamlit as st
import torch
import numpy as np
import os
import cv2
import tempfile
import pandas as pd
import altair as alt
import re
from transformers import AutoProcessor, AutoModelForVideoClassification
import matplotlib.pyplot as plt

# Set page configuration
st.set_page_config(layout="wide", page_title="Action Recognition")

# Sidebar
st.sidebar.write("## Upload and Process Video 🎥")
uploaded_file = st.sidebar.file_uploader("Upload a video file:", type=["mp4", "avi", "mov"])

# Sidebar Information
with st.sidebar.expander("ℹ️ Video Guidelines"):
    st.write("""
    - Supported formats: MP4, AVI, MOV
    - Ensure the video contains clear actions for better predictions
    """)

def download_model_if_needed(save_path):
    if not os.path.exists(save_path):
        st.info("Downloading model from Google Drive...")
        # This is your actual shared model file ID from Google Drive
        file_id = "1yegsjiRVRtXpLfaIpisNPSX6B931sbTG"
        url = f"https://drive.google.com/uc?id={file_id}"
        gdown.download(url, save_path, quiet=False)
        st.success("✅ Model downloaded successfully!")

@st.cache_resource
def load_model():
    model_path = "/home/urk24cs1210/24KIDS416/src/training/final_best_timesformer_model.pth"
    download_model_if_needed(model_path)

    model = AutoModelForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400")
    model.classifier = torch.nn.Linear(model.config.hidden_size, 25)  # adjust to match your dataset class count
    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
    extractor = AutoFeatureExtractor.from_pretrained("facebook/timesformer-base-finetuned-k400")
    return model, extractor

# Function to extract frames from a video
def extract_frames_from_video(video_path, output_folder, num_frames=8):
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_interval = max(total_frames // num_frames, 1)

    frame_count = 0
    saved_frames = 0
    while cap.isOpened() and saved_frames < num_frames:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % frame_interval == 0:
            frame_path = os.path.join(output_folder, f"frame_{saved_frames + 1:04d}.jpg")
            frame = cv2.resize(frame, (224, 224))
            cv2.imwrite(frame_path, frame)
            saved_frames += 1
        frame_count += 1
    cap.release()

# Main Layout
st.write("## Action Recognition App")
st.write("Upload a video to predict the action using a pre-trained model.")

# Introduction
st.write("""
This app allows you to upload a video, converts it into frames, and predicts the action using a pre-trained model.
We use **TimeSformer**, a state-of-the-art video transformer model, which processes video frames as a sequence of images and captures temporal relationships to predict actions effectively.
Experience seamless action recognition with visualizations and confidence scores.
""")

# Two-column layout
col1, col2 = st.columns(2)

if uploaded_file:
    with tempfile.TemporaryDirectory() as temp_dir:
        video_path = os.path.join(temp_dir, uploaded_file.name)
        with open(video_path, "wb") as f:
            f.write(uploaded_file.read())

        # Display the uploaded video
        col1.write("### Uploaded Video")
        col1.video(video_path)

        # Extract frames from the video
        st.info("Extracting frames from the video...")
        extract_frames_from_video(video_path, temp_dir, num_frames=8)
        folder_path = temp_dir

        # Process the extracted frames
        image_files = sorted([f for f in os.listdir(folder_path) if f.endswith(".jpg")])[:8]
        frames = []

        for img_name in image_files:
            img_path = os.path.join(folder_path, img_name)
            frame = cv2.imread(img_path)
            frames.append(frame)

        if len(frames) < 8:
            st.warning("The video must contain enough frames to extract 8 frames.")
        else:
            # Use processor instead of extractor
            inputs = processor([frames], return_tensors="pt")

            with torch.no_grad():
                outputs = model(**inputs)
                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
                top_prob, top_index = torch.max(probs, dim=-1)

            # Display the single top prediction
            col2.write("### Predicted Action")
            action_label = model.config.id2label[top_index.item()]
            confidence = top_prob.item() * 100
            col2.markdown(
                f"""
                <div style="background-color: #f9f9f9; padding: 10px; border-radius: 10px; box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1); margin-bottom: 10px;">
                    <h2 style="font-size: 24px; color: #4CAF50;">{action_label}</h2>
                    <p style="font-size: 16px; color: #777;">Confidence: {confidence:.2f}%</p>
                </div>
                """,
                unsafe_allow_html=True,
            )

            # Generate heatmaps for visualization
            heatmaps = []
            for idx, frame in enumerate(frames):
                # Create a random heatmap for demonstration purposes
                heatmap = np.zeros((224, 224), dtype=np.uint8)
                center_x, center_y = 112 + (idx * 10) % 50, 112 + (idx * 10) % 50
                cv2.circle(heatmap, (center_x, center_y), 50, (255), -1)
                heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
                overlay = cv2.addWeighted(frame, 0.6, heatmap, 0.4, 0)
                heatmaps.append(overlay)

            # Display the frames and heatmaps
            st.write("### Heatmap Visualization")
            fig, axes = plt.subplots(2, 8, figsize=(20, 5))
            for i in range(8):
                if i < len(frames):
                    axes[0, i].imshow(cv2.cvtColor(frames[i], cv2.COLOR_BGR2RGB))
                    axes[0, i].axis("off")
                if i < len(heatmaps):
                    axes[1, i].imshow(cv2.cvtColor(heatmaps[i], cv2.COLOR_BGR2RGB))
                    axes[1, i].axis("off")
            st.pyplot(fig)

# Training Loss Curve
st.write("## Training Loss Curve")
try:
    losses = []
    log_file_path = "logs/training.log"  # Path to the training log file

    # Check if the log file exists
    if os.path.exists(log_file_path):
        with open(log_file_path, "r") as file:
            for line in file:
                # Extract loss values using a regular expression
                match = re.search(r"Loss: ([0-9.]+)", line)
                if match:
                    losses.append(float(match.group(1)))

        # If losses are found, plot the training loss curve
        if losses:
            df = pd.DataFrame({"Epoch": range(1, len(losses) + 1), "Loss": losses})
            chart = (
                alt.Chart(df)
                .mark_line(point=True)
                .encode(
                    x=alt.X("Epoch:Q", title="Epochs"),
                    y=alt.Y("Loss:Q", title="Loss"),
                    tooltip=["Epoch", "Loss"],
                )
                .properties(title="Training Loss Curve", width=800, height=400)
                .interactive()
            )
            st.altair_chart(chart, use_container_width=True)
        else:
            st.warning("The training log file is empty or does not contain valid data.")
    else:
        st.warning(f"Training log file not found. Please ensure the file exists at '{log_file_path}'.")
except Exception as e:
    st.warning(f"An error occurred while reading the training log: {str(e)}")