Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,69 +1,32 @@
 import streamlit as st
-import cv2
-import torch
-import numpy as np
-from PIL import Image
+from video_utils import extract_frames, annotate_video
+from model import predict_fault
+import os
 
-# Load pre-trained YOLOv5 model (change to your model if custom)
-@st.cache_resource
-def load_model():
-    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # YOLOv5 small model (you can use your custom model here)
-    return model
+st.title("🔍 Solar Panel Fault Detection from Video")
+st.write("Upload a drone video of solar panels to detect faults: **cracked**, **dusted**, **shaded**, **overheated**.")
 
-# Extract frames from the video
-def extract_frames(video_path):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert to RGB
-        frames.append(rgb_frame)
-    cap.release()
-    return frames
+uploaded_file = st.file_uploader("Upload a solar panel video", type=["mp4", "avi"])
 
-# Process frames with YOLOv5 model
-def process_frames_with_model(frames, model):
-    results = []
-    for frame in frames:
-        img = Image.fromarray(frame)
-        results.append(model(img))  # Use YOLOv5 to detect objects
-    return results
+if uploaded_file is not None:
+    # Save uploaded video temporarily
+    video_path = os.path.join("temp_video.mp4")
+    with open(video_path, "wb") as f:
+        f.write(uploaded_file.read())
 
-# Display results
-def display_results(results):
-    for idx, result in enumerate(results):
-        st.subheader(f"Frame {idx+1}")
-        st.image(result.render()[0])  # Render and display frame with bounding boxes
-        st.write("Detected Faults:")
-        st.write(result.pandas().xywh)  # Display predictions as a table
+    st.video(video_path)
 
-# Main function for Streamlit interface
-def main():
-    st.title("Solar Fault Detection - Thermal Image Analysis")
+    st.write("🔎 Processing video...")
+    frames = extract_frames(video_path, interval=30)  # every 30th frame (~1/sec)
 
-    st.sidebar.header("Upload Video")
-    uploaded_video = st.sidebar.file_uploader("Choose a video file", type=["mp4", "mov"])
+    predictions = {}
+    for frame_idx, image in frames:
+        label = predict_fault(image)
+        predictions[frame_idx] = label
+        st.image(image, caption=f"Frame {frame_idx}: {label}", use_column_width=True)
 
-    if uploaded_video is not None:
-        video_path = "/mnt/data/uploaded_video.mp4"
-        with open(video_path, "wb") as f:
-            f.write(uploaded_video.read())
-        st.video(uploaded_video)  # Display video preview
+    st.write("🎞️ Generating annotated video...")
+    output_video_path = annotate_video(video_path, predictions)
 
-        # Load YOLOv5 model (or your custom model)
-        model = load_model()
-
-        # Extract frames from the video
-        frames = extract_frames(video_path)
-
-        # Process frames with the YOLOv5 model
-        results = process_frames_with_model(frames, model)
-
-        # Display results on the Streamlit app
-        display_results(results)
-
-# Run the app
-if __name__ == "__main__":
-    main()
+    st.video(output_video_path)
+    st.success("✅ Detection complete.")