src/streamlit_app.py

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

st.title("Pole Fault Detection")
st.write("Upload a video (.mp4) or image (.jpg, .png) to detect pole faults.")

# File uploader
uploaded_file = st.file_uploader("Choose a file", type=["mp4", "jpg", "png"])

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

    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 = []

        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")

        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")

    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")