Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,117 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import cv2
+import pandas as pd
+import tempfile
+import os
+from services.video_service import get_video_frame
+from services.detection_service import detect_objects
+from services.fault_service import detect_pole_faults
+from services.salesforce_dispatcher import send_to_salesforce
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+st.title("Pole Fault Detection")
+st.write("Upload a video (.mp4) or image (.jpg, .png) to detect pole faults.")
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+# File uploader
+uploaded_file = st.file_uploader("Choose a file", type=["mp4", "jpg", "png"])
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+if uploaded_file is not None:
+    # Save uploaded file to temporary location
+    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{uploaded_file.name.split('.')[-1]}") as tmp_file:
+        tmp_file.write(uploaded_file.read())
+        file_path = tmp_file.name
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+    if uploaded_file.type.startswith("video"):
+        # Process video
+        st.subheader("Processed Video")
+        video_placeholder = st.empty()
+        fault_table = st.empty()
+        frame_gen = get_video_frame(file_path)
+        output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        cap = cv2.VideoCapture(file_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        faults_list = []
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
+        for frame in frame_gen:
+            # Save frame temporarily for detection
+            cv2.imwrite("temp.jpg", frame)
+            detections = detect_objects("temp.jpg")
+            faults = detect_pole_faults("temp.jpg")
+            alert_payload = {
+                "detections": detections,
+                "faults": bool(faults),
+                "fault_details": faults
+            }
+            # Send to Salesforce
+            send_to_salesforce(alert_payload)
+            faults_list.extend(faults)
+            # Annotate frame (basic text overlay for faults)
+            for fault in faults:
+                cv2.putText(
+                    frame,
+                    f"{fault['fault_type']} ({fault['confidence']:.2f})",
+                    (50, 50),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1,
+                    (0, 0, 255),
+                    2
+                )
+            out.write(frame)
+            # Display frame in Streamlit
+            video_placeholder.image(frame, channels="BGR")
 
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
+        out.release()
+        cap.release()
+        # Display final video
+        st.video(output_path)
+        # Display faults table
+        if faults_list:
+            df = pd.DataFrame(faults_list)
+            fault_table.subheader("Detected Faults")
+            fault_table.dataframe(df)
+        else:
+            fault_table.write("No faults detected.")
+        # Clean up
+        os.remove(file_path)
+        os.remove(output_path)
+        os.remove("temp.jpg")
 
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+    else:
+        # Process image
+        st.subheader("Processed Image")
+        image = cv2.imread(file_path)
+        cv2.imwrite("temp.jpg", image)
+        detections = detect_objects("temp.jpg")
+        faults = detect_pole_faults("temp.jpg")
+        alert_payload = {
+            "detections": detections,
+            "faults": bool(faults),
+            "fault_details": faults
+        }
+        # Send to Salesforce
+        send_to_salesforce(alert_payload)
+        # Annotate image (basic text overlay for faults)
+        for fault in faults:
+            cv2.putText(
+                image,
+                f"{fault['fault_type']} ({fault['confidence']:.2f})",
+                (50, 50),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 0, 255),
+                2
+            )
+        st.image(image, channels="BGR", caption="Processed Image")
+        # Display faults table
+        if faults:
+            st.subheader("Detected Faults")
+            df = pd.DataFrame(faults)
+            st.dataframe(df)
+        else:
+            st.write("No faults detected.")
+        # Clean up
+        os.remove(file_path)
+        os.remove("temp.jpg")