JO-7/Action_Recognition

1import streamlit as st2import torch3import numpy as np4import os5import cv26import tempfile7import pandas as pd8import altair as alt9import re10from transformers import AutoProcessor, AutoModelForVideoClassification11import matplotlib.pyplot as plt12 13# Set page configuration14st.set_page_config(layout="wide", page_title="Action Recognition")15 16# Sidebar17st.sidebar.write("## Upload and Process Video 🎥")18uploaded_file = st.sidebar.file_uploader("Upload a video file:", type=["mp4", "avi", "mov"])19 20# Sidebar Information21with st.sidebar.expander("ℹ️ Video Guidelines"):22    st.write("""23    - Supported formats: MP4, AVI, MOV24    - Ensure the video contains clear actions for better predictions25    """)26 27def download_model_if_needed(save_path):28    if not os.path.exists(save_path):29        st.info("Downloading model from Google Drive...")30        # This is your actual shared model file ID from Google Drive31        file_id = "1yegsjiRVRtXpLfaIpisNPSX6B931sbTG"32        url = f"https://drive.google.com/uc?id={file_id}"33        gdown.download(url, save_path, quiet=False)34        st.success("✅ Model downloaded successfully!")35 36@st.cache_resource37def load_model():38    model_path = "/home/urk24cs1210/24KIDS416/src/training/final_best_timesformer_model.pth"39    download_model_if_needed(model_path)40 41    model = AutoModelForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400")42    model.classifier = torch.nn.Linear(model.config.hidden_size, 25)  # adjust to match your dataset class count43    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))44    extractor = AutoFeatureExtractor.from_pretrained("facebook/timesformer-base-finetuned-k400")45    return model, extractor46 47# Function to extract frames from a video48def extract_frames_from_video(video_path, output_folder, num_frames=8):49    cap = cv2.VideoCapture(video_path)50    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))51    frame_interval = max(total_frames // num_frames, 1)52 53    frame_count = 054    saved_frames = 055    while cap.isOpened() and saved_frames < num_frames:56        ret, frame = cap.read()57        if not ret:58            break59        if frame_count % frame_interval == 0:60            frame_path = os.path.join(output_folder, f"frame_{saved_frames + 1:04d}.jpg")61            frame = cv2.resize(frame, (224, 224))62            cv2.imwrite(frame_path, frame)63            saved_frames += 164        frame_count += 165    cap.release()66 67# Main Layout68st.write("## Action Recognition App")69st.write("Upload a video to predict the action using a pre-trained model.")70 71# Introduction72st.write("""73This app allows you to upload a video, converts it into frames, and predicts the action using a pre-trained model.74We 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.75Experience seamless action recognition with visualizations and confidence scores.76""")77 78# Two-column layout79col1, col2 = st.columns(2)80 81if uploaded_file:82    with tempfile.TemporaryDirectory() as temp_dir:83        video_path = os.path.join(temp_dir, uploaded_file.name)84        with open(video_path, "wb") as f:85            f.write(uploaded_file.read())86 87        # Display the uploaded video88        col1.write("### Uploaded Video")89        col1.video(video_path)90 91        # Extract frames from the video92        st.info("Extracting frames from the video...")93        extract_frames_from_video(video_path, temp_dir, num_frames=8)94        folder_path = temp_dir95 96        # Process the extracted frames97        image_files = sorted([f for f in os.listdir(folder_path) if f.endswith(".jpg")])[:8]98        frames = []99 100        for img_name in image_files:101            img_path = os.path.join(folder_path, img_name)102            frame = cv2.imread(img_path)103            frames.append(frame)104 105        if len(frames) < 8:106            st.warning("The video must contain enough frames to extract 8 frames.")107        else:108            # Use processor instead of extractor109            inputs = processor([frames], return_tensors="pt")110 111            with torch.no_grad():112                outputs = model(**inputs)113                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)114                top_prob, top_index = torch.max(probs, dim=-1)115 116            # Display the single top prediction117            col2.write("### Predicted Action")118            action_label = model.config.id2label[top_index.item()]119            confidence = top_prob.item() * 100120            col2.markdown(121                f"""122                <div style="background-color: #f9f9f9; padding: 10px; border-radius: 10px; box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1); margin-bottom: 10px;">123                    <h2 style="font-size: 24px; color: #4CAF50;">{action_label}</h2>124                    <p style="font-size: 16px; color: #777;">Confidence: {confidence:.2f}%</p>125                </div>126                """,127                unsafe_allow_html=True,128            )129 130            # Generate heatmaps for visualization131            heatmaps = []132            for idx, frame in enumerate(frames):133                # Create a random heatmap for demonstration purposes134                heatmap = np.zeros((224, 224), dtype=np.uint8)135                center_x, center_y = 112 + (idx * 10) % 50, 112 + (idx * 10) % 50136                cv2.circle(heatmap, (center_x, center_y), 50, (255), -1)137                heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)138                overlay = cv2.addWeighted(frame, 0.6, heatmap, 0.4, 0)139                heatmaps.append(overlay)140 141            # Display the frames and heatmaps142            st.write("### Heatmap Visualization")143            fig, axes = plt.subplots(2, 8, figsize=(20, 5))144            for i in range(8):145                if i < len(frames):146                    axes[0, i].imshow(cv2.cvtColor(frames[i], cv2.COLOR_BGR2RGB))147                    axes[0, i].axis("off")148                if i < len(heatmaps):149                    axes[1, i].imshow(cv2.cvtColor(heatmaps[i], cv2.COLOR_BGR2RGB))150                    axes[1, i].axis("off")151            st.pyplot(fig)152 153# Training Loss Curve154st.write("## Training Loss Curve")155try:156    losses = []157    log_file_path = "logs/training.log"  # Path to the training log file158 159    # Check if the log file exists160    if os.path.exists(log_file_path):161        with open(log_file_path, "r") as file:162            for line in file:163                # Extract loss values using a regular expression164                match = re.search(r"Loss: ([0-9.]+)", line)165                if match:166                    losses.append(float(match.group(1)))167 168        # If losses are found, plot the training loss curve169        if losses:170            df = pd.DataFrame({"Epoch": range(1, len(losses) + 1), "Loss": losses})171            chart = (172                alt.Chart(df)173                .mark_line(point=True)174                .encode(175                    x=alt.X("Epoch:Q", title="Epochs"),176                    y=alt.Y("Loss:Q", title="Loss"),177                    tooltip=["Epoch", "Loss"],178                )179                .properties(title="Training Loss Curve", width=800, height=400)180                .interactive()181            )182            st.altair_chart(chart, use_container_width=True)183        else:184            st.warning("The training log file is empty or does not contain valid data.")185    else:186        st.warning(f"Training log file not found. Please ensure the file exists at '{log_file_path}'.")187except Exception as e:188    st.warning(f"An error occurred while reading the training log: {str(e)}")189