Spaces:
Runtime error
Runtime error
import streamlit as st | |
import torch | |
import cv2 | |
import numpy as np | |
from PIL import Image | |
import tempfile | |
import os | |
# Load the model from local path | |
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt') | |
def detect_fall(image): | |
results = model(image) | |
return results | |
def main(): | |
st.title("Fall Detection Application") | |
option = st.sidebar.selectbox("Choose input type", ("Upload Image/Video", "Use Camera")) | |
if option == "Upload Image/Video": | |
uploaded_file = st.file_uploader("Upload Image or Video", type=['jpg', 'jpeg', 'png', 'mp4', 'avi', 'mov']) | |
if uploaded_file is not None: | |
if uploaded_file.type.startswith('image'): | |
file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8) | |
image = cv2.imdecode(file_bytes, 1) | |
st.image(image, caption='Uploaded Image.', use_column_width=True) | |
results = detect_fall(image) | |
st.image(np.squeeze(results.render()), caption='Processed Image.', use_column_width=True) | |
elif uploaded_file.type.startswith('video'): | |
tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
tfile.write(uploaded_file.read()) | |
tfile.close() | |
st.write(f'Temporary file path: {tfile.name}') | |
vidcap = cv2.VideoCapture(tfile.name) | |
stframe = st.empty() | |
while vidcap.isOpened(): | |
success, frame = vidcap.read() | |
if not success: | |
break | |
results = detect_fall(frame) | |
processed_frame = np.squeeze(results.render()) | |
stframe.image(processed_frame, channels="BGR") | |
vidcap.release() | |
os.remove(tfile.name) | |
elif option == "Use Camera": | |
stframe = st.empty() | |
cap = cv2.VideoCapture(0) | |
while True: | |
ret, frame = cap.read() | |
if not ret: | |
break | |
results = detect_fall(frame) | |
processed_frame = np.squeeze(results.render()) | |
stframe.image(processed_frame, channels="BGR") | |
cap.release() | |
if __name__ == '__main__': | |
main() | |