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"""

{action_label}

Confidence: {confidence:.2f}%

""", 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)}")