Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,26 +1,19 @@
+import streamlit as st
+import torch
 import cv2
 import numpy as np
-import streamlit as st
 from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
 
-st.title("Live Fault Detection in Utility Poles")
+# Load the YOLOv5 model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='model/best.pt', force_reload=True)
 
-class FaultDetector(VideoTransformerBase):
+st.title("Utility Pole Fault Detection")
+
+class VideoTransformer(VideoTransformerBase):
     def transform(self, frame):
         img = frame.to_ndarray(format="bgr24")
+        results = model(img)
+        annotated_frame = np.squeeze(results.render())
+        return annotated_frame
 
-        # Convert to grayscale
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-
-        # Apply Gaussian blur
-        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
-
-        # Perform Canny edge detection
-        edges = cv2.Canny(blurred, 50, 150)
-
-        # Convert edges to 3-channel image
-        edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
-
-        return edges_colored
-
-webrtc_streamer(key="fault_detection", video_transformer_factory=FaultDetector)
+webrtc_streamer(key="live", video_transformer_factory=VideoTransformer)