File size: 4,124 Bytes
9f77b8d
562f13d
 
 
 
 
 
 
 
9f77b8d
562f13d
 
9f77b8d
562f13d
 
9f77b8d
562f13d
 
 
 
 
9f77b8d
562f13d
 
 
 
 
 
 
 
 
 
 
 
 
 
9f77b8d
562f13d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f77b8d
562f13d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f77b8d
562f13d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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")