Update app.py

app.py

@@ -19,16 +19,17 @@ def load_model():
 
 model = load_model()
 
-# Fault detection based on intensity simulation
-def detect_faults(frame, results):
+# Fault detection with frame & location tracking
+def detect_faults(frame, results, frame_number):
     faults = {"Thermal Fault": False, "Dust Fault": False, "Power Generation Fault": False}
+    fault_locations = []
     annotated_frame = frame.copy()
 
     for result in results:
         boxes = result.boxes
         for box in boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
-            conf = box.conf[0]
+            conf = float(box.conf[0])
             cls = int(box.cls[0])
 
             roi = frame[y1:y2, x1:x2]
@@ -47,13 +48,20 @@ def detect_faults(frame, results):
             else:
                 continue
 
-            # Transparent overlay
+            # Track location
+            fault_locations.append({
+                "frame": frame_number,
+                "label": label,
+                "confidence": round(conf, 2),
+                "intensity": round(mean_intensity, 2),
+                "box": (x1, y1, x2, y2)
+            })
+
+            # Annotate
             overlay = annotated_frame.copy()
             alpha = 0.3
             cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1)
             cv2.addWeighted(overlay, alpha, annotated_frame, 1 - alpha, 0, annotated_frame)
-
-            # Border and label
             cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
             cv2.putText(annotated_frame, f"{label} ({mean_intensity:.1f})", (x1, y1 - 10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
@@ -61,14 +69,14 @@ def detect_faults(frame, results):
     if faults["Thermal Fault"] or faults["Dust Fault"]:
         faults["Power Generation Fault"] = True
 
-    return annotated_frame, faults
+    return annotated_frame, faults, fault_locations
 
 # Video processing
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         st.error("Failed to open video.")
-        return None, None
+        return None, None, None
 
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -80,6 +88,7 @@ def process_video(video_path):
 
     frame_count = 0
     video_faults = {"Thermal Fault": False, "Dust Fault": False, "Power Generation Fault": False}
+    all_fault_locations = []
     process_every_n_frames = fps  # 1 frame per second
 
     with st.spinner("Processing video..."):
@@ -90,16 +99,17 @@ def process_video(video_path):
                 break
 
             if frame_count % process_every_n_frames == 0:
-                resized = cv2.resize(frame, (640, 480))  # Resize for faster inference
+                resized = cv2.resize(frame, (640, 480))
                 frame_rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
                 results = model(frame_rgb, verbose=False)
 
-                annotated_frame, faults = detect_faults(frame, results)
+                annotated_frame, faults, locations = detect_faults(frame, results, frame_count)
+                all_fault_locations.extend(locations)
 
                 for fault in video_faults:
                     video_faults[fault] |= faults[fault]
             else:
-                annotated_frame = frame  # Keep original for non-processed frames
+                annotated_frame = frame
 
             out.write(annotated_frame)
             frame_count += 1
@@ -107,7 +117,7 @@ def process_video(video_path):
 
     cap.release()
     out.release()
-    return output_path, video_faults
+    return output_path, video_faults, all_fault_locations
 
 # File uploader
 uploaded_file = st.file_uploader("Upload a thermal video", type=["mp4"])
@@ -119,7 +129,7 @@ if uploaded_file:
 
     st.video(temp_input_path)
 
-    output_path, video_faults = process_video(temp_input_path)
+    output_path, video_faults, fault_locations = process_video(temp_input_path)
 
     if output_path:
         st.subheader("Detection Results")
@@ -140,6 +150,19 @@ if uploaded_file:
         else:
             st.success("No faults detected. The system seems to be functioning properly.")
 
+        # Display fault locations
+        if fault_locations:
+            st.subheader("📍 Fault Locations in Video")
+            st.dataframe([
+                {
+                    "Frame": loc["frame"],
+                    "Fault Type": loc["label"],
+                    "Confidence": loc["confidence"],
+                    "Intensity": loc["intensity"],
+                    "Box": f"{loc['box']}"
+                } for loc in fault_locations
+            ])
+
         os.unlink(output_path)
     os.unlink(temp_input_path)