Spaces:

Medvira
/

Workouts

Sleeping

Update app.py

cb20121 verified over 2 years ago

2.75 kB

import streamlit as st
import cv2
import tempfile
import numpy as np
from ultralytics import YOLO, solutions

# Load the YOLOv8 model
model = YOLO("yolov8n-pose.pt")

# Streamlit App
st.title("Workout Monitoring App")
st.write("Upload a video to monitor your ab workout.")

uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi"])

if uploaded_file is not None:
    # Save the uploaded video to a temporary file
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(uploaded_file.read())
    tfile.close()

    # Load the video with OpenCV
    cap = cv2.VideoCapture(tfile.name)
    assert cap.isOpened(), "Error reading video file"
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Initialize AIGym object
    gym_object = solutions.AIGym(
        line_thickness=2,
        view_img=False,  # Set to False since we are using Streamlit to display
        pose_type="abworkout",  # Use 'abworkout' as the pose type
        kpts_to_check=[6, 8, 10],
    )

    # List to store processed frames
    processed_frames = []

    # Streamlit progress bar
    progress_bar = st.progress(0)
    frame_count = 0

    # Process the video frame by frame
    st.write("Analyzing video... Please wait.")
    while cap.isOpened():
        success, im0 = cap.read()
        if not success:
            break

        results = model.track(im0, verbose=False)  # Tracking recommended
        im0 = gym_object.start_counting(im0, results)

        # Resize the frame to 320x320
        im0_resized = cv2.resize(im0, (1024, 1024))

        # Append the processed frame to the list
        processed_frames.append(im0_resized)

        # Update progress bar
        frame_count += 1
        progress_bar.progress(frame_count / total_frames)

    cap.release()
    cv2.destroyAllWindows()

    # Create a temporary file to save the processed video
    output_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    output_video_path = output_video_file.name

    # Write the processed frames to a video file
    video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (1024, 1024))
    for frame in processed_frames:
        video_writer.write(frame)
    video_writer.release()

    # Display the processed video in Streamlit
    st.write("Analysis complete. Displaying processed video:")
    st.video(output_video_path)

    # Provide a download link for the processed video
    st.write("Download the processed video:")
    with open(output_video_path, "rb") as video_file:
        st.download_button("Download", video_file, "workouts.mp4")

	import streamlit as st
	import cv2
	import tempfile
	import numpy as np
	from ultralytics import YOLO, solutions

	# Load the YOLOv8 model
	model = YOLO("yolov8n-pose.pt")

	# Streamlit App
	st.title("Workout Monitoring App")
	st.write("Upload a video to monitor your ab workout.")

	uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi"])

	if uploaded_file is not None:
	# Save the uploaded video to a temporary file
	tfile = tempfile.NamedTemporaryFile(delete=False)
	tfile.write(uploaded_file.read())
	tfile.close()

	# Load the video with OpenCV
	cap = cv2.VideoCapture(tfile.name)
	assert cap.isOpened(), "Error reading video file"
	w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
	total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

	# Initialize AIGym object
	gym_object = solutions.AIGym(
	line_thickness=2,
	view_img=False, # Set to False since we are using Streamlit to display
	pose_type="abworkout", # Use 'abworkout' as the pose type
	kpts_to_check=[6, 8, 10],
	)

	# List to store processed frames
	processed_frames = []

	# Streamlit progress bar
	progress_bar = st.progress(0)
	frame_count = 0

	# Process the video frame by frame
	st.write("Analyzing video... Please wait.")
	while cap.isOpened():
	success, im0 = cap.read()
	if not success:
	break

	results = model.track(im0, verbose=False) # Tracking recommended
	im0 = gym_object.start_counting(im0, results)

	# Resize the frame to 320x320
	im0_resized = cv2.resize(im0, (1024, 1024))

	# Append the processed frame to the list
	processed_frames.append(im0_resized)

	# Update progress bar
	frame_count += 1
	progress_bar.progress(frame_count / total_frames)

	cap.release()
	cv2.destroyAllWindows()

	# Create a temporary file to save the processed video
	output_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
	output_video_path = output_video_file.name

	# Write the processed frames to a video file
	video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (1024, 1024))
	for frame in processed_frames:
	video_writer.write(frame)
	video_writer.release()

	# Display the processed video in Streamlit
	st.write("Analysis complete. Displaying processed video:")
	st.video(output_video_path)

	# Provide a download link for the processed video
	st.write("Download the processed video:")
	with open(output_video_path, "rb") as video_file:
	st.download_button("Download", video_file, "workouts.mp4")