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

# Set a writable cache directory for PyTorch
torch_cache_dir = os.path.join(os.getcwd(), 'torch_cache')
os.makedirs(torch_cache_dir, exist_ok=True)
os.environ['TORCH_HOME'] = torch_cache_dir

# 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)