import streamlit as st
import torch
import cv2
import numpy as np
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase

# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='model/best.pt', force_reload=True)

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

webrtc_streamer(key="live", video_transformer_factory=VideoTransformer)