Rename src/streamlit_app.py to app.py

app.py

@@ -0,0 +1,37 @@
+import streamlit as st
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+
+# Load your pre-trained model
+model = torch.load('model/your_model_file.pt')
+model.eval()
+
+# Define image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
+
+st.title("VIEP: Utility Pole Fault Detection")
+
+uploaded_file = st.file_uploader("Upload an image of a utility pole", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file).convert('RGB')
+    st.image(image, caption='Uploaded Image', use_column_width=True)
+
+    # Preprocess the image
+    input_tensor = transform(image).unsqueeze(0)
+
+    # Perform inference
+    with torch.no_grad():
+        output = model(input_tensor)
+        _, predicted = torch.max(output, 1)
+
+    # Map the prediction to class names
+    classes = ['No Fault', 'Fault Detected']
+    st.write(f"Prediction: {classes[predicted.item()]}")

src/streamlit_app.py

@@ -1,24 +0,0 @@
-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)