Spaces:
Build error
Build error
Update src/streamlit_app.py
Browse files- 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 |
-
|
|
|
7 |
|
8 |
-
|
|
|
|
|
9 |
def transform(self, frame):
|
10 |
img = frame.to_ndarray(format="bgr24")
|
|
|
|
|
|
|
11 |
|
12 |
-
|
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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|