Sanjayraju30 commited on
Commit
563c1d2
·
verified ·
1 Parent(s): 3472efb

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +11 -18
src/streamlit_app.py CHANGED
@@ -1,26 +1,19 @@
 
 
1
  import cv2
2
  import numpy as np
3
- import streamlit as st
4
  from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
5
 
6
- st.title("Live Fault Detection in Utility Poles")
 
7
 
8
- class FaultDetector(VideoTransformerBase):
 
 
9
  def transform(self, frame):
10
  img = frame.to_ndarray(format="bgr24")
 
 
 
11
 
12
- # Convert to grayscale
13
- gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
14
-
15
- # Apply Gaussian blur
16
- blurred = cv2.GaussianBlur(gray, (5, 5), 0)
17
-
18
- # Perform Canny edge detection
19
- edges = cv2.Canny(blurred, 50, 150)
20
-
21
- # Convert edges to 3-channel image
22
- edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
23
-
24
- return edges_colored
25
-
26
- webrtc_streamer(key="fault_detection", video_transformer_factory=FaultDetector)
 
1
+ import streamlit as st
2
+ import torch
3
  import cv2
4
  import numpy as np
 
5
  from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
6
 
7
+ # Load the YOLOv5 model
8
+ model = torch.hub.load('ultralytics/yolov5', 'custom', path='model/best.pt', force_reload=True)
9
 
10
+ st.title("Utility Pole Fault Detection")
11
+
12
+ class VideoTransformer(VideoTransformerBase):
13
  def transform(self, frame):
14
  img = frame.to_ndarray(format="bgr24")
15
+ results = model(img)
16
+ annotated_frame = np.squeeze(results.render())
17
+ return annotated_frame
18
 
19
+ webrtc_streamer(key="live", video_transformer_factory=VideoTransformer)