Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,16 +7,21 @@ import pandas as pd
 from PIL import Image
 import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
+from huggingface_hub import login
 
-# === Load Hugging Face Model using default cache (handles permissions safely) ===
-processor = DetrImageProcessor.from_pretrained("NaveenKumar5/Solar_panel_fault_detection")
-model = DetrForObjectDetection.from_pretrained("NaveenKumar5/Solar_panel_fault_detection")
+# Get token from environment (set as secret in HF Spaces)
+HUGGINGFACE_TOKEN = os.environ.get("HF_TOKEN")
+login(token=HUGGINGFACE_TOKEN)
+
+# === Load Model and Processor ===
+processor = DetrImageProcessor.from_pretrained("NaveenKumar5/Solar_panel_fault_detection", use_auth_token=True)
+model = DetrForObjectDetection.from_pretrained("NaveenKumar5/Solar_panel_fault_detection", use_auth_token=True)
 model.eval()
 
-# === Streamlit App Setup ===
+# === Streamlit UI Setup ===
 st.set_page_config(page_title="Solar Panel Fault Detection", layout="wide")
-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.")
+st.title("🔍 Solar Panel Fault Detection (Hugging Face Model)")
+st.write("Upload a thermal video (MP4). Faults will be detected using your Hugging Face DETR model.")
 
 # === Fault Detection Function ===
 def detect_faults(frame, frame_idx, fps):
@@ -34,7 +39,7 @@ def detect_faults(frame, frame_idx, fps):
         x1, y1, x2, y2 = map(int, box.tolist())
         conf = score.item()
         label_id = label.item()
-        label_name = f"class_{label_id}"
+        label_name = f"class_{label_id}"  # Customize if label mapping is available
 
         # Draw on frame
         color = (0, 0, 255)
@@ -49,9 +54,10 @@ def detect_faults(frame, frame_idx, fps):
             "Confidence": round(conf, 2),
             "Box": f"({x1},{y1},{x2},{y2})"
         })
+
     return frame, faults
 
-# === Video Processing Function ===
+# === Process Uploaded Video ===
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     fps = int(cap.get(cv2.CAP_PROP_FPS))
@@ -94,21 +100,22 @@ if uploaded_file:
     with open(temp_input_path, "wb") as f:
         f.write(uploaded_file.read())
 
-    output_path, log = process_video(temp_input_path)
-
-    st.subheader("🧪 Processed Output")
-    st.video(output_path)
-
-    if log:
-        df = pd.DataFrame(log)
-        st.write("### 📊 Detected Faults Table")
-        st.dataframe(df)
-        st.download_button("📥 Download Fault Log CSV", convert_df(df), "fault_log.csv", "text/csv")
-    else:
-        st.success("✅ No faults detected.")
-
-    os.unlink(temp_input_path)
-    os.unlink(output_path)
+    try:
+        output_path, log = process_video(temp_input_path)
+
+        st.subheader("🧪 Processed Output")
+        st.video(output_path)
+
+        if log:
+            df = pd.DataFrame(log)
+            st.write("### 📊 Detected Faults Table")
+            st.dataframe(df)
+            st.download_button("📥 Download Fault Log CSV", convert_df(df), "fault_log.csv", "text/csv")
+        else:
+            st.success("✅ No faults detected.")
+    finally:
+        os.unlink(temp_input_path)
+        os.unlink(output_path)
 
 st.markdown("---")
 st.caption("Built with Streamlit + Hugging Face DETR + OpenCV")