Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |
| ) | |