Spaces:
Runtime error
Runtime error
File size: 3,343 Bytes
8d18193 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import streamlit as st
from PIL import Image
import os
from utils.yolo_processor import YOLOProcessor
import tempfile
import numpy as np
import base64
processed_image = None
processed_video_path = None
def detect_fall(image, model_path):
model = YOLOProcessor(model_path)
result_image = model.detect_fall(image)
return result_image
def main():
global processed_image, processed_video_path
st.title("Fall Detection with YOLO")
st.markdown("---")
option = st.sidebar.selectbox("Choose an option", ["Image", "Video"])
if option == "Image":
st.subheader("Upload Image")
uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image', use_column_width=True)
st.markdown("---")
st.subheader("Detecting Fall...")
if processed_image is None: # Process the image only if it hasn't been processed before
with st.spinner('Detecting fall...'):
processed_image = detect_fall(image, "assets/best.pt")
st.image(processed_image, caption='Result', use_column_width=True)
# Download button for the result image
if st.button('Download Result Image'):
download_image(processed_image, filename='result_image.png')
elif option == "Video":
st.subheader("Upload Video")
uploaded_file = st.file_uploader("Choose a video", type=["mp4"])
if uploaded_file is not None:
st.markdown("---")
st.subheader("Processing and Detecting Fall...")
temp_dir = tempfile.TemporaryDirectory()
temp_file_path = os.path.join(temp_dir.name, "uploaded_video.mp4")
with open(temp_file_path, "wb") as f:
f.write(uploaded_file.read())
output_path = os.path.join(temp_dir.name, "processed_video.mp4")
if processed_video_path is None:
with st.spinner('Processing and detecting fall...'):
yolo_processor = YOLOProcessor("assets/best.pt")
yolo_processor.process_video(temp_file_path, output_path)
processed_video_path = output_path
st.subheader("Result Video")
st.video(processed_video_path)
if st.button('Download Result Video'):
download_file(processed_video_path, filename='processed_video.mp4')
temp_dir.cleanup()
def download_image(image, filename):
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
image.save(filename)
with open(filename, "rb") as f:
image_bytes = f.read()
b64 = base64.b64encode(image_bytes).decode()
href = f'<a href="data:image/png;base64,{b64}" download="{filename}">Click here to download {filename}</a>'
st.markdown(href, unsafe_allow_html=True)
def download_file(file_path, filename):
with open(file_path, 'rb') as f:
data = f.read()
b64 = base64.b64encode(data).decode()
href = f'<a href="data:file/mp4;base64,{b64}" download="{filename}">Click here to download {filename}</a>'
st.markdown(href, unsafe_allow_html=True)
if __name__ == "__main__":
main()
|