Spaces:
Sleeping
Sleeping
| import matplotlib | |
| matplotlib.use("Agg") | |
| # Libraries | |
| import streamlit as st | |
| # Utility | |
| import time | |
| import io | |
| import argparse | |
| import sys | |
| import os.path | |
| import subprocess | |
| import tempfile | |
| import logging | |
| # Viz | |
| import altair as alt | |
| import av | |
| # ML | |
| import numpy as np | |
| import pandas as pd | |
| import cv2 as cv | |
| from PIL import Image, ImageOps | |
| from tqdm import tqdm | |
| # Custom | |
| import inference | |
| from app_utils import * | |
| def run_app(): | |
| # Options | |
| st.set_option("deprecation.showfileUploaderEncoding", False) | |
| # Counter variable | |
| views_counter = 0 | |
| # Increment the counter on page/view access | |
| views_counter += 1 | |
| # Display the counter value | |
| st.write("Number of Views:", views_counter) | |
| # App documentation | |
| st.title("MIT Count Fish Counter") | |
| st.text("Upload a video file to detect and count fish") | |
| video_url = "yolo2_out_py.mp4" # Replace with the actual video URL or file path | |
| video_bytes = open(video_url, "rb").read() | |
| col1, col2 = st.columns(2) | |
| ## Col1 ######################################### | |
| with col1: | |
| ## Initial visualizations | |
| # Historical values | |
| st.altair_chart( | |
| plot_historical_data(pd.read_csv("herring_count_all.csv")), | |
| use_container_width=True, | |
| ) | |
| # Locations | |
| st.subheader("Map of Fishery Locations") | |
| st.map( | |
| pd.DataFrame( | |
| np.random.randn(5, 2) / [50, 50] + [42.41, -71.38], | |
| columns=["lat", "lon"], | |
| ) | |
| ) | |
| ## Col2 ######################################### | |
| with col2: | |
| # Replace the loading message with the video | |
| st.subheader("Example of processed video") | |
| st.video(video_bytes) | |
| st.subheader("Upload your own video...") | |
| # Initialize | |
| img_types = ["jpg", "png", "jpeg"] | |
| video_types = ["mp4", "avi"] | |
| uploaded_file = st.file_uploader( | |
| "Select an image or video file...", type=img_types + video_types | |
| ) | |
| # Display uploaded file | |
| if uploaded_file is not None: | |
| if str(uploaded_file.type).split("/")[-1] in img_types: | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded image", use_column_width=True) | |
| # TBD: Inference code to run and display for single image | |
| elif str(uploaded_file.type).split("/")[-1] in video_types: | |
| # Display uploaded file | |
| st.video(uploaded_file) | |
| # Convert streamlit video object to OpenCV format to run inferences | |
| tfile = tempfile.NamedTemporaryFile(delete=False) | |
| tfile.write(uploaded_file.read()) | |
| vf = cv.VideoCapture(tfile.name) | |
| # Run inference | |
| with st.spinner("Running inference..."): | |
| frames, counts, timestamps = inference.main(vf) | |
| logging.info("INFO: Completed running inference on frames") | |
| st.balloons() | |
| # Convert OpenCV Numpy frames in-memory to IO Bytes for streamlit | |
| streamlit_video_file = frames_to_video(frames=frames, fps=11) | |
| st.video(streamlit_video_file) # Show processed video\ | |
| st.download_button( | |
| label="Download processed video", | |
| data=streamlit_video_file, | |
| mime="mp4", | |
| file_name="processed_video.mp4", | |
| ) | |
| df_counts_time = pd.DataFrame( | |
| data={"fish_count": counts, "timestamps": timestamps[1:]} | |
| ) | |
| st.altair_chart( | |
| plot_count_date(dataframe=df_counts_time), | |
| use_container_width=True, | |
| ) | |
| else: | |
| st.write("No file uploaded") | |
| if __name__ == "__main__": | |
| run_app() | |