Update inference.py

inference.py

@@ -41,6 +41,68 @@ def load_historical(fpath):
 
 st.set_page_config(layout="wide")
 
+
+def process_uploaded_file():
+    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,
+            )
+
+    else:
+        st.write("No file uploaded")
+
+
 # Define the main function to run the Streamlit app
 def run_app():
     # Set Streamlit options
@@ -57,92 +119,36 @@ def run_app():
     # 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,
-        )
-
-        # 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"],
+    main_tab, upload_tab = st.tabs(["Analysis", "Upload video for analysis"])
+
+
+
+    with main_tab:
+        # 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,
+            )
+
+            # 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"],
+                ),use_container_width=True
             )
-        )
-
-    ## 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,
-                )
-
-        else:
-            st.write("No file uploaded")
+        with col2:
+            # Display example processed video
+            st.subheader("Example of processed video")
+            st.video(video_bytes)
+    with upload_tab:
+        process_uploaded_file()
 
 # Run the app if the script is executed directly
 if __name__ == "__main__":