Spaces:
Sleeping
Sleeping
File size: 2,925 Bytes
74cd228 |
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 |
"""
Description: Functions that directly support the Streamlit app
"""
import pandas as pd
import altair as alt
import io
import av
from tqdm import tqdm
import numpy as np
import logging
def frames_to_video(frames=None, fps=12):
"""
Convert frames to video for Streamlit
Args:
frames: frame from cv2.VideoCapture as numpy. E.g. frame.astype(np.uint8)
fps: Frames per second. Usefull if inference video is compressed to slow down for analysis
"""
height, width, layers = frames[0].shape # grab info from first frame
output_memory_file = io.BytesIO() # Create BytesIO "in memory file".
output = av.open(
output_memory_file, "w", format="mp4"
) # Open "in memory file" as MP4 video output
stream = output.add_stream(
"h264", str(fps)
) # Add H.264 video stream to the MP4 container, with framerate = fps.
stream.width = width # Set frame width
stream.height = height # Set frame height
stream.pix_fmt = "yuv420p" # NOTE: yuv444p doesn't work on mac. Select yuv444p pixel format (better quality than default yuv420p).
stream.options = {
"crf": "17"
} # Select low crf for high quality (the price is larger file size).
# Iterate the created images, encode and write to MP4 memory file.
logging.info("INFO: Encoding frames and writing to MP4 format.")
for frame in tqdm(frames):
frame = av.VideoFrame.from_ndarray(frame.astype(np.uint8), format="bgr24")
packet = stream.encode(frame) # Encode video frame
output.mux(
packet
) # "Mux" the encoded frame (add the encoded frame to MP4 file).
packet = stream.encode(None) # Flush the encoder
output.mux(packet)
output.close()
output_memory_file.seek(0)
return output_memory_file
def plot_historical_data(dataframe):
"""Returns altair plot of historical counts to be rendered on main dashboard."""
dataframe["Date"] = pd.to_datetime(dataframe["Date"])
s = (
dataframe.resample(rule="D", on="Date")["Count"].sum().reset_index()
) # Resample on day
return (
alt.Chart(s, title="Historical Video Counts of Herring")
.mark_bar()
.transform_window(
# The field to average
rolling_mean="mean(Count)",
# The number of values before and after the current value to include.
frame=[-9, 0],
)
.encode(x="Date", y="Count", tooltip=["Count", "Date"])
.interactive()
)
def plot_count_date(dataframe):
"""Plots counts vs relative time for uploaded video."""
dataframe["seconds"] = dataframe["timestamps"] / 1000
dataframe["class"] = "Herring" # TBD: Hard-coded for now
return (
alt.Chart(dataframe, title="Processed video detected fish")
.mark_line()
.encode(x="seconds", y="fish_count", color="class")
.interactive()
)
|