Update app.py

app.py

@@ -3,15 +3,19 @@ import gradio as gr
 import cv2
 import time
 import json
+import random
 import logging
 import matplotlib.pyplot as plt
 import shutil
 from datetime import datetime
+from collections import Counter
 from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 
+# Suppress Ultralytics warning by setting a writable config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
 
+# Import service modules
 try:
     from services.video_service import get_next_video_frame, reset_video_index, preload_video, release_video
     from services.metrics_service import update_metrics
@@ -20,15 +24,21 @@ try:
     from services.thermal_service import process_thermal
     from services.map_service import generate_map
     from services.under_construction.earthwork_detection import process_earthwork
-    from services.under_construction.culvert_detection import process_culverts
-    from services.under_construction.bridge_pier_detection import process_bridge_piers
+    from services.under_construction.culvert_check import process_culverts
+    from services.under_construction.bridge_pier_check import process_bridge_piers
 except ImportError as e:
     print(f"Failed to import service modules: {str(e)}")
     logging.error(f"Import error: {str(e)}")
     exit(1)
 
-logging.basicConfig(filename="app.log", level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+# Configure logging
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
+# Global variables
 paused: bool = False
 frame_rate: float = 0.3
 frame_count: int = 0
@@ -44,25 +54,22 @@ active_service: Optional[str] = None
 is_video: bool = True
 static_image: Optional[np.ndarray] = None
 enabled_services: List[str] = []
-video_writer: Optional[cv2.VideoWriter] = None
-output_video_path: str = ""
 
+# Constants
 DEFAULT_VIDEO_PATH = "sample.mp4"
 TEMP_IMAGE_PATH = os.path.abspath("temp.jpg")
 CAPTURED_FRAMES_DIR = os.path.abspath("captured_frames")
 OUTPUT_DIR = os.path.abspath("outputs")
 TEMP_MEDIA_DIR = os.path.abspath("temp_media")
 
+# Ensure directories exist with write permissions
 for directory in [CAPTURED_FRAMES_DIR, OUTPUT_DIR, TEMP_MEDIA_DIR]:
     os.makedirs(directory, exist_ok=True)
     os.chmod(directory, 0o777)
 
 def initialize_media(media_file: Optional[Any] = None) -> str:
-    global media_loaded, is_video, static_image, log_entries, frame_count, video_writer, output_video_path
+    global media_loaded, is_video, static_image, log_entries, frame_count
     release_video()
-    if video_writer is not None:
-        video_writer.release()
-        video_writer = None
     static_image = None
     frame_count = 0
 
@@ -103,16 +110,6 @@ def initialize_media(media_file: Optional[Any] = None) -> str:
     try:
         if file_extension in (".mp4", ".avi"):
             is_video = True
-            cap = cv2.VideoCapture(media_path)
-            fps = cap.get(cv2.CAP_PROP_FPS)
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            cap.release()
-
-            output_video_path = os.path.join(OUTPUT_DIR, f"processed_under_construction_{int(time.time())}.mp4")
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-            video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
-
             preload_video(media_path)
             media_loaded = True
             status = f"Successfully loaded video: {media_path}"
@@ -184,12 +181,11 @@ def monitor_feed() -> Tuple[
     str,
     List[str],
     Optional[str],
-    Optional[str],
     Optional[str]
 ]:
     global paused, frame_count, last_frame, last_metrics, last_timestamp
-    global gps_coordinates, detected_issues, media_loaded, video_writer
-    global is_video, static_image, enabled_services, output_video_path
+    global gps_coordinates, detected_issues, media_loaded
+    global is_video, static_image, enabled_services
 
     if not media_loaded:
         log_entries.append("Cannot start processing: Media not loaded successfully.")
@@ -200,7 +196,6 @@ def monitor_feed() -> Tuple[
             "\n".join(log_entries[-10:]),
             detected_issues,
             None,
-            None,
             None
         )
 
@@ -210,13 +205,14 @@ def monitor_feed() -> Tuple[
     else:
         max_retries = 3
         start_time = time.time()
+        frame = None
         for attempt in range(max_retries):
             try:
                 if is_video:
                     frame = get_next_video_frame()
-                    if frame is None:
-                        log_entries.append(f"Frame retrieval failed on attempt {attempt + 1}, resetting video.")
-                        logging.warning(f"Frame retrieval failed on attempt {attempt + 1}, resetting video.")
+                    if frame is None or frame.size == 0:
+                        log_entries.append(f"Empty frame on attempt {attempt + 1}, resetting video.")
+                        logging.warning(f"Empty frame on attempt {attempt + 1}, resetting video.")
                         reset_video_index()
                         continue
                     break
@@ -233,10 +229,9 @@ def monitor_feed() -> Tuple[
                         "\n".join(log_entries[-10:]),
                         detected_issues,
                         None,
-                        None,
                         None
                     )
-        else:
+        if frame is None:
             log_entries.append("Failed to retrieve frame after maximum retries.")
             logging.error("Failed to retrieve frame after maximum retries.")
             return (
@@ -245,10 +240,23 @@ def monitor_feed() -> Tuple[
                 "\n".join(log_entries[-10:]),
                 detected_issues,
                 None,
-                None,
                 None
             )
 
+        # Skip frame if processing is too slow
+        processing_time = time.time() - start_time
+        if is_video and processing_time > frame_rate:
+            log_entries.append(f"Processing too slow ({processing_time*1000:.0f}ms > {frame_rate*1000:.0f}ms), skipping frame.")
+            logging.warning(f"Processing too slow, skipping frame.")
+            return (
+                last_frame[:, :, ::-1] if last_frame is not None else None,
+                json.dumps(last_metrics, indent=2),
+                "\n".join(log_entries[-10:]),
+                detected_issues,
+                generate_line_chart(),
+                map_path if 'map_path' in locals() else None
+            )
+
         detection_frame = cv2.resize(frame, (512, 320))
         all_detected_items: List[Dict[str, Any]] = []
         shadow_issue = False
@@ -258,8 +266,8 @@ def monitor_feed() -> Tuple[
             if "under_construction" in enabled_services:
                 earthwork_dets, detection_frame = process_earthwork(detection_frame)
                 culvert_dets, detection_frame = process_culverts(detection_frame)
-                pier_dets, detection_frame = process_bridge_piers(detection_frame)
-                all_detected_items.extend(earthwork_dets + culvert_dets + pier_dets)
+                bridge_pier_dets, detection_frame = process_bridge_piers(detection_frame)
+                all_detected_items.extend(earthwork_dets + culvert_dets + bridge_pier_dets)
 
             try:
                 cv2.imwrite(TEMP_IMAGE_PATH, detection_frame)
@@ -294,22 +302,20 @@ def monitor_feed() -> Tuple[
                 if not box:
                     continue
                 x_min, y_min, x_max, y_max = box
+                label = item.get("label", "")
                 dtype = item.get("type", "")
 
                 if dtype == "earthwork":
-                    color = (255, 105, 180)  # Pink
-                    label = "Earthwork"
+                    color = (255, 20, 147)  # Bright Pink
                 elif dtype == "culvert":
-                    color = (0, 128, 128)  # Teal
-                    label = "Culvert"
+                    color = (0, 255, 255)  # Bright Teal
                 elif dtype == "bridge_pier":
-                    color = (255, 127, 127)  # Coral
-                    label = "Bridge Pier"
+                    color = (255, 99, 71)  # Bright Coral
                 else:
                     continue
 
-                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 3)
-                (text_w, text_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 4)
+                (text_w, text_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
                 label_background = frame[y_min - text_h - 15:y_min - 5, x_min:x_min + text_w + 10]
                 if label_background.size > 0:
                     overlay = label_background.copy()
@@ -317,7 +323,7 @@ def monitor_feed() -> Tuple[
                     alpha = 0.5
                     cv2.addWeighted(overlay, alpha, label_background, 1 - alpha, 0, label_background)
                 cv2.putText(frame, label, (x_min + 5, y_min - 10),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
 
             try:
                 cv2.imwrite(TEMP_IMAGE_PATH, frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
@@ -331,7 +337,7 @@ def monitor_feed() -> Tuple[
             all_detected_items = []
 
         metrics = update_metrics(all_detected_items)
-        gps_coord = [17.385044, 78.486671]  # Fixed GPS for consistency
+        gps_coord = [17.385044 + random.uniform(-0.001, 0.001), 78.486671 + frame_count * 0.0001]
         gps_coordinates.append(gps_coord)
 
         for item in all_detected_items:
@@ -345,10 +351,9 @@ def monitor_feed() -> Tuple[
                 if not success:
                     raise RuntimeError(f"Failed to save captured frame: {captured_frame_path}")
                 for item in all_detected_items:
-                    if item.get("type") in ["earthwork", "culvert", "bridge_pier"]:
-                        detected_issues.append(captured_frame_path)
-                        if len(detected_issues) > 100:
-                            detected_issues.pop(0)
+                    detected_issues.append(captured_frame_path)
+                    if len(detected_issues) > 100:
+                        detected_issues.pop(0)
             except Exception as e:
                 log_entries.append(f"Error saving captured frame: {str(e)}")
                 logging.error(f"Error saving captured frame: {str(e)}")
@@ -369,12 +374,14 @@ def monitor_feed() -> Tuple[
             log_entries.append(f"Salesforce Dispatch Error: {str(e)}")
             logging.error(f"Salesforce dispatch error: {str(e)}")
 
-        if is_video and video_writer is not None:
-            try:
-                video_writer.write(frame)
-            except Exception as e:
-                log_entries.append(f"Error writing frame to video: {str(e)}")
-                logging.error(f"Error writing frame to video: {str(e)}")
+        try:
+            frame_path = os.path.join(OUTPUT_DIR, f"frame_{frame_count:04d}.jpg")
+            success = cv2.imwrite(frame_path, frame, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
+            if not success:
+                raise RuntimeError(f"Failed to save output frame: {frame_path}")
+        except Exception as e:
+            log_entries.append(f"Error saving output frame: {str(e)}")
+            logging.error(f"Error saving output frame: {str(e)}")
 
         frame_count += 1
         last_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
@@ -383,16 +390,16 @@ def monitor_feed() -> Tuple[
 
         earthwork_detected = len([item for item in all_detected_items if item.get("type") == "earthwork"])
         culvert_detected = len([item for item in all_detected_items if item.get("type") == "culvert"])
-        pier_detected = len([item for item in all_detected_items if item.get("type") == "bridge_pier"])
-        detected_counts.append(earthwork_detected + culvert_detected + pier_detected)
+        bridge_pier_detected = len([item for item in all_detected_items if item.get("type") == "bridge_pier"])
+        detected_counts.append(earthwork_detected + culvert_detected + bridge_pier_detected)
 
         processing_time = time.time() - start_time
         detection_summary = {
             "timestamp": last_timestamp,
             "frame": frame_count,
-            "earthwork": earthwork_detected,
+            "earthworks": earthwork_detected,
             "culverts": culvert_detected,
-            "bridge_piers": pier_detected,
+            "bridge_piers": bridge_pier_detected,
             "gps": gps_coord,
             "processing_time_ms": processing_time * 1000
         }
@@ -412,30 +419,22 @@ def monitor_feed() -> Tuple[
     map_items = [item for item in last_metrics.get("items", []) if item.get("type") in ["earthwork", "culvert", "bridge_pier"]]
     map_path = generate_map(gps_coordinates[-5:], map_items)
 
-    download_link = output_video_path if is_video and os.path.exists(output_video_path) else None
-
     return (
         frame[:, :, ::-1],
         json.dumps(last_metrics, indent=2),
         "\n".join(log_entries[-10:]),
         detected_issues,
         generate_line_chart(),
-        map_path,
-        download_link
+        map_path
     )
 
-def finalize_video() -> str:
-    global video_writer
-    if video_writer is not None:
-        video_writer.release()
-        video_writer = None
-        log_entries.append("Processed video saved successfully.")
-        logging.info("Processed video saved successfully.")
-        return "Video processing completed."
-    return "No video to finalize."
-
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="amber")) as app:
-    gr.Markdown("# 🏗️ Under Construction Inspection Dashboard\nMonitor construction sites in real-time using drone footage or static images.")
+    gr.Markdown(
+        """
+        # 🛠️ Under Construction Inspection Dashboard
+        Monitor under construction elements in real-time using drone footage or static images.
+        """
+    )
 
     with gr.Row():
         with gr.Column(scale=3):
@@ -476,21 +475,26 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="amber")
         chart_output = gr.Image(label="Detection Trend")
         map_output = gr.Image(label="Issue Locations Map")
 
-    with gr.Row():
-        video_download = gr.File(label="Download Processed Video", interactive=False)
-
     with gr.Row():
         pause_btn = gr.Button("⏸️ Pause", variant="secondary")
         resume_btn = gr.Button("▶️ Resume", variant="primary")
-        stop_btn = gr.Button("⏹️ Stop and Save Video", variant="secondary")
         frame_slider = gr.Slider(0.05, 1.0, value=0.3, label="Frame Interval (seconds)", step=0.05)
 
     gr.HTML("""
     <style>
-        body { background-color: #FFDAB9 !important; }
-        #live-feed { border: 2px solid #FF8C00; border-radius: 10px; }
-        .gr-button-primary { background-color: #FF8C00 !important; }
-        .gr-button-secondary { background-color: #FF6347 !important; }
+        body {
+            background-color: #FFDAB9 !important;
+        }
+        #live-feed {
+            border: 2px solid #FF8C00;
+            border-radius: 10px;
+        }
+        .gr-button-primary {
+            background-color: #FF8C00 !important;
+        }
+        .gr-button-secondary {
+            background-color: #FF6347 !important;
+        }
     </style>
     """)
 
@@ -510,7 +514,11 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="amber")
 
     media_status.value = initialize_media()
 
-    load_button.click(initialize_media, inputs=[media_input], outputs=[media_status])
+    load_button.click(
+        initialize_media,
+        inputs=[media_input],
+        outputs=[media_status]
+    )
 
     def update_toggles(uc_val: bool) -> Tuple[str, str]:
         active, status_message = set_active_service(uc_val)
@@ -518,26 +526,26 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="amber")
         return uc_status_val, status_message
 
     uc_toggle.change(update_toggles, inputs=[uc_toggle], outputs=[uc_status, status_text])
+
     pause_btn.click(toggle_pause, outputs=status_text)
     resume_btn.click(toggle_resume, outputs=status_text)
-    stop_btn.click(finalize_video, outputs=status_text)
     frame_slider.change(set_frame_rate, inputs=[frame_slider])
 
     def streaming_loop():
         while True:
             if not media_loaded:
-                yield None, json.dumps({"error": "Media not loaded. Please upload a video or image file."}, indent=2), "\n".join(log_entries[-10:]), detected_issues, None, None, None
+                yield None, json.dumps({"error": "Media not loaded. Please upload a video or image file."}, indent=2), "\n".join(log_entries[-10:]), detected_issues, None, None
             else:
-                frame, metrics, logs, issues, chart, map_path, download_link = monitor_feed()
+                frame, metrics, logs, issues, chart, map_path = monitor_feed()
                 if frame is None:
-                    yield None, metrics, logs, issues, chart, map_path, download_link
+                    yield None, metrics, logs, issues, chart, map_path
                 else:
-                    yield frame, metrics, logs, issues, chart, map_path, download_link
+                    yield frame, metrics, logs, issues, chart, map_path
                 if not is_video:
                     break
             time.sleep(frame_rate)
 
-    app.load(streaming_loop, outputs=[media_output, metrics_output, logs_output, issue_images, chart_output, map_output, video_download])
+    app.load(streaming_loop, outputs=[media_output, metrics_output, logs_output, issue_images, chart_output, map_output])
 
 if __name__ == "__main__":
     app.launch(share=True)