Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,3 @@
-# src/streamlit_app.py
 import streamlit as st
 import cv2
 import numpy as np
@@ -9,15 +8,24 @@ from PIL import Image
 import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
 
-# === Hugging Face DETR Configuration ===
-processor = DetrImageProcessor.from_pretrained("NaveenKumar5/Solar_panel_fault_detection")
-model = DetrForObjectDetection.from_pretrained("NaveenKumar5/Solar_panel_fault_detection")
+# === Ensure cache directory for model download ===
+os.makedirs("./model_cache", exist_ok=True)
+
+# === Load Hugging Face Model ===
+processor = DetrImageProcessor.from_pretrained(
+    "NaveenKumar5/Solar_panel_fault_detection",
+    cache_dir="./model_cache"
+)
+model = DetrForObjectDetection.from_pretrained(
+    "NaveenKumar5/Solar_panel_fault_detection",
+    cache_dir="./model_cache"
+)
 model.eval()
 
-# === Streamlit App Configuration ===
+# === Streamlit App ===
 st.set_page_config(page_title="Solar Panel Fault Detection", layout="wide")
-st.title("\U0001F50D Solar Panel Fault Detection (Hugging Face DETR)")
-st.write("Upload a thermal video (MP4). Faults will be detected using your custom DETR model.")
+st.title("🔍 Solar Panel Fault Detection (DETR - Hugging Face)")
+st.write("Upload a thermal video (MP4). Faults will be detected using a DETR model from Hugging Face.")
 
 # === Fault Detection Function ===
 def detect_faults(frame, frame_idx, fps):
@@ -37,6 +45,7 @@ def detect_faults(frame, frame_idx, fps):
         label_id = label.item()
         label_name = f"class_{label_id}"
 
+        # Draw on frame
         color = (0, 0, 255)
         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
         cv2.putText(frame, f"{label_name} ({conf:.2f})", (x1, y1 - 5),
@@ -51,7 +60,7 @@ def detect_faults(frame, frame_idx, fps):
         })
     return frame, faults
 
-# === Video Processing Function ===
+# === Video Processing ===
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     fps = int(cap.get(cv2.CAP_PROP_FPS))
@@ -81,12 +90,12 @@ def process_video(video_path):
     writer.release()
     return output_path, fault_log
 
-# === CSV Conversion ===
+# === CSV Download Helper ===
 def convert_df(df):
     return df.to_csv(index=False).encode('utf-8')
 
-# === Streamlit UI ===
-uploaded_file = st.file_uploader("\U0001F4E4 Upload thermal video", type=["mp4"])
+# === Streamlit Interface ===
+uploaded_file = st.file_uploader("📤 Upload thermal video", type=["mp4"])
 if uploaded_file:
     st.video(uploaded_file)
 
@@ -96,16 +105,16 @@ if uploaded_file:
 
     output_path, log = process_video(temp_input_path)
 
-    st.subheader("\U0001F9EA Processed Output")
+    st.subheader("🧪 Processed Output")
     st.video(output_path)
 
     if log:
         df = pd.DataFrame(log)
-        st.write("### \U0001F4CA Detected Faults Table")
+        st.write("### 📊 Detected Faults Table")
         st.dataframe(df)
-        st.download_button("\U0001F4E5 Download Fault Log CSV", convert_df(df), "fault_log.csv", "text/csv")
+        st.download_button("📥 Download Fault Log CSV", convert_df(df), "fault_log.csv", "text/csv")
     else:
-        st.success("\u2705 No faults detected.")
+        st.success("✅ No faults detected.")
 
     os.unlink(temp_input_path)
     os.unlink(output_path)