Update inference.py

inference.py

@@ -1,168 +1,149 @@
-# This code is written at BigVision LLC. It is based on the OpenCV project. It is subject to the license terms in the LICENSE file found in this distribution and at http://opencv.org/license.html
+# Import necessary libraries
+import matplotlib
 
-# Usage example:  python3 object_detection_yolo.py --video=run.mp4 --device 'cpu'
-#                 python3 object_detection_yolo.py --video=run.mp4 --device 'gpu'
-#                 python3 object_detection_yolo.py --image=bird.jpg --device 'cpu'
-#                 python3 object_detection_yolo.py --image=bird.jpg --device 'gpu'
-# python3 test2.py --image=/Users/apowell/Downloads/HerringInTrap.JPG --device 'cpu'
-# python3 test2.py --video=/Users/apowell/Downloads/sampleFull.avi --device 'cpu'
-# sampleFull.avi
+# Use Agg backend for Matplotlib
+matplotlib.use("Agg")
 
-import cv2 as cv
+# Libraries for the app
+import streamlit as st
+import time
+import io
 import argparse
 import sys
-import numpy as np
 import os.path
-import os
-import matplotlib
-import streamlit as st
-
-matplotlib.use("Agg")
-from inference_utils import *
-from PIL import Image, ImageOps
+import subprocess
+import tempfile
 import logging
 
-# Custom
-from centroidtracker import CentroidTracker
-
-# Set default static images for testing while working locally
-DEFAULT_IMAGE = "/Users/apowell/Downloads/HerringInTrap.JPG"
-DEFAULT_VIDEO = "/Users/apowell/Downloads/sampleFull.avi"
-YOUTUBE = "https://www.youtube.com/watch?v=CbB7vl_HUbU&ab_channel=AustinPowell"
-
-
-def main(input_file=None, is_image=False, device="cpu"):
-    """
-    Run main inference script. Returns annotated frames from inference and counts of fish.
-
-    Args:
-        - input_file: image or video file input from OpenCV
-        - is_image: Binary denoting single image
-        - device: CPU or GPU processing
-    """
-    ## Initialize the parameters
-    # Confidence threshold
-    conf_threshold = 0.5
-    # Non-maximum suppression threshold (maximum bounding box)
-    nms_threshold = 0.05
-    input_width = 416  # Width of network's input image
-    input_height = 416  # Height of network's input image
-
-    # Generic name assignment for output file
-    outputFile = "yolo2_out_py.mp4"
-    # Load class name
-    classes = "Herring"
-    # Give the configuration and weight files for the model and load the network using them.
-    modelConfiguration = "herring.cfg"
-    modelWeights = "herring_final.weights"
-
-    # Centroid tracker to Id specific objects (NOTE: This is temporary and not fully tested)
-    tracker = CentroidTracker(maxDisappeared=80, maxDistance=90)
-
-    # Process inputs
-    if (
-        type(input_file) == cv.VideoCapture
-    ):  # Video objects passed from something like Streamlit
-        cap = input_file
-    elif type(input_file) == str:  # For local uploads
-        cap = cv.VideoCapture(input_file)
-        logging.info("INFO: Loading file locally: {}".format(input_file))
-    else:
-        sys.exit(
-            "Input file is of type {} and not solved for.".format(type(input_file))
-        )
+# Visualization libraries
+import altair as alt
+import av
 
-    net = cv.dnn.readNetFromDarknet(modelConfiguration, modelWeights)
-
-    # Get the video writer initialized to save the output video
-    vid_writer = cv.VideoWriter(
-        outputFile,
-        cv.VideoWriter_fourcc("M", "J", "P", "G"),
-        30,
-        (
-            round(cap.get(cv.CAP_PROP_FRAME_WIDTH)),
-            round(cap.get(cv.CAP_PROP_FRAME_HEIGHT)),
-        ),
-    )
-
-    total_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
-    video_fps = cap.get(cv.CAP_PROP_FPS)
-    logging.info(
-        "INFO: Starting inference process on video frames. Total: {}, fps: {}".format(
-            total_frames, video_fps
-        )
-    )
-    timestamps = [cap.get(cv.CAP_PROP_POS_MSEC)]  # Timestamp for frame
-    calc_timestamps = [0.0]  # Relative timestamps to first timestamp
-    saved_frames = []  # Save CV2 frames
-    count_list = []
-    while cap.isOpened():
-
-        # Get frame from the video
-        hasFrame, frame = cap.read()
-
-        # Stop the program if reached end of video
-        if not hasFrame:
-            print("Done processing !!!")
-            print("Output file is stored as ", outputFile)
-            # Release device
-            cap.release()
-            break
-
-        # Create a 4D blob from a frame.
-        blob = cv.dnn.blobFromImage(
-            frame, 1 / 255, (input_width, input_height), [0, 0, 0], 1, crop=False
+# Machine Learning and Image Processing libraries
+import numpy as np
+import pandas as pd
+import cv2 as cv
+from PIL import Image, ImageOps
+from tqdm import tqdm
+
+# Custom modules
+import inference
+from app_utils import *
+
+@st.cache_data
+def load_video(video_url):
+    video_bytes = open(video_url, "rb").read()
+    return video_bytes
+
+@st.cache_data
+def load_historical(fpath):
+    return pd.read_csv(fpath)
+
+st.set_page_config(layout="wide")
+
+# Define the main function to run the Streamlit app
+def run_app():
+    # Set Streamlit options
+    st.set_option("deprecation.showfileUploaderEncoding", False)
+    
+    # App title and description
+    st.title("MIT Count Fish Counter")
+    st.text("Upload a video file to detect and count fish")
+    
+    # Example video URL or file path (replace with actual video URL or file path)
+    video_url = "yolo2_out_py.mp4"
+    video_bytes = load_video(video_url)
+
+    # Load historical herring
+    df_historical_herring = load_historical(fpath="herring_count_all.csv")
+
+    tab1, map_tab = st.tabs(["📈 Chart", "Map of Fishery Locations"])
+
+    # Create two columns for layout
+    col1, col2 = st.columns(2)
+
+    ## Col1 #########################################
+    with col1:
+        ## Initial visualizations
+        # Plot historical data
+        st.altair_chart(
+            plot_historical_data(df_historical_herring),
+            use_container_width=True,
         )
 
-        # Sets the input to the network
-        net.setInput(blob)
-
-        # Runs the forward pass to get output of the output layers
-        outs = net.forward(get_outputs_names(net=net))
-
-        # Remove the bounding boxes with low confidence
-        counts = postprocess(
-            frame=frame,
-            outs=outs,
-            tracker=tracker,
-            conf_threshold=conf_threshold,
-            nms_threshold=nms_threshold,
-            classes=classes,
+        # Display map of fishery 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"],
+            )
         )
-        count_list.append(counts)
 
-        # Put efficiency information. The function getPerfProfile returns the overall time for inference(t) and the timings for each of the layers(in layersTimes)
-        t, _ = net.getPerfProfile()
-        label = "Inference time: %.2f ms" % (t * 1000.0 / cv.getTickFrequency())
-        cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
+    ## Col2 #########################################
+    with col2:
+        # Display example processed video
+        st.subheader("Example of processed video")
+        st.video(video_bytes)
+        st.subheader("Upload your own video...")
+
+        # Initialize accepted file types for upload
+        img_types = ["jpg", "png", "jpeg"]
+        video_types = ["mp4", "avi"]
+
+        # Allow user to upload an image or video file
+        uploaded_file = st.file_uploader("Select an image or video file...", type=img_types + video_types)
+
+        # Display the uploaded file
+        if uploaded_file is not None:
+            if str(uploaded_file.type).split("/")[-1] in img_types:
+                # Display uploaded image
+                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 video
+                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 on the uploaded video
+                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)
+
+                # Show processed video and provide download button
+                st.video(streamlit_video_file)
+                st.download_button(
+                    label="Download processed video",
+                    data=streamlit_video_file,
+                    mime="mp4",
+                    file_name="processed_video.mp4",
+                )
+
+                # Create dataframe for fish counts and timestamps
+                df_counts_time = pd.DataFrame(
+                    data={"fish_count": counts, "timestamps": timestamps[1:]}
+                )
+
+                # Display fish count vs. timestamp chart
+                st.altair_chart(
+                    plot_count_date(dataframe=df_counts_time),
+                    use_container_width=True,
+                )
 
-        # Save frame
-        saved_frames.append(frame.astype(np.uint8))  # )
-
-        # Write the frame with the detection boxes
-        if is_image:
-            cv.imwrite(outputFile, frame.astype(np.uint8))
         else:
-            vid_writer.write(frame.astype(np.uint8))
-
-        timestamps.append(cap.get(cv.CAP_PROP_POS_MSEC))
-        calc_timestamps.append(calc_timestamps[-1] + 1000 / video_fps)
-    # Calculate time difference for different timestamps
-    time_diffs = [
-        abs(ts - cts) for i, (ts, cts) in enumerate(zip(timestamps, calc_timestamps))
-    ]
-
-    with open("your_file.csv", "w") as f:
-        for i in range(len(count_list)):
-            f.write(
-                f"{count_list[i]}, {time_diffs[i+1]}, {timestamps[i]}, {calc_timestamps[i]}\n"
-            )
-
-    return saved_frames, count_list, timestamps
-
+            st.write("No file uploaded")
 
+# Run the app if the script is executed directly
 if __name__ == "__main__":
-
-    # Script below to enable running pure inference from command line
-    file_path = "/Users/apowell/Downloads/2_2018-04-27_15-50-53.mp4"
-    saved_frames, counts, timestamps = main(input_file=file_path)
+    run_app()