Update app.py

app.py

@@ -58,16 +58,17 @@ paused: bool = False
 frame_rate: float = 0.3
 frame_count: int = 0
 log_entries: List[str] = []
-crack_counts: List[int] = []
-crack_severity_all: List[str] = []
+detected_counts: List[int] = []
 last_frame: Optional[np.ndarray] = None
 last_metrics: Dict[str, Any] = {}
 last_timestamp: str = ""
-last_detected_cracks: List[str] = []
-last_detected_holes: List[str] = []
+detected_plants: List[str] = []  # For plants and missing patches
+detected_issues: List[str] = []  # For cracks, holes, and missing patches
 gps_coordinates: List[List[float]] = []
-video_loaded: bool = False
+media_loaded: bool = False
 active_service: Optional[str] = None
+is_video: bool = True
+static_image: Optional[np.ndarray] = None
 
 # Constants
 DEFAULT_VIDEO_PATH = "sample.mp4"
@@ -77,28 +78,47 @@ OUTPUT_DIR = "outputs"
 os.makedirs(CAPTURED_FRAMES_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-def initialize_video(video_file: Optional[Any] = None) -> str:
-    global video_loaded, log_entries
+def initialize_media(media_file: Optional[Any] = None) -> str:
+    global media_loaded, is_video, static_image, log_entries
     release_video()
-    video_path = DEFAULT_VIDEO_PATH if video_file is None else video_file.name
+    static_image = None
+    media_path = DEFAULT_VIDEO_PATH if media_file is None else media_file.name
 
-    if not os.path.exists(video_path):
-        status = f"Error: Video file '{video_path}' not found."
+    if not os.path.exists(media_path):
+        status = f"Error: Media file '{media_path}' not found."
         log_entries.append(status)
         logging.error(status)
-        video_loaded = False
+        media_loaded = False
         return status
 
     try:
-        preload_video(video_path)
-        video_loaded = True
-        status = f"Successfully loaded video: {video_path}"
+        # Determine if the file is a video or image
+        if media_path.lower().endswith((".mp4", ".avi")):
+            is_video = True
+            preload_video(media_path)
+            media_loaded = True
+            status = f"Successfully loaded video: {media_path}"
+        elif media_path.lower().endswith((".jpg", ".jpeg", ".png")):
+            is_video = False
+            static_image = cv2.imread(media_path)
+            if static_image is None:
+                raise RuntimeError(f"Failed to load image: {media_path}")
+            static_image = cv2.resize(static_image, (320, 240))
+            media_loaded = True
+            status = f"Successfully loaded image: {media_path}"
+        else:
+            media_loaded = False
+            status = "Error: Unsupported file format. Use .mp4, .avi, .jpg, .jpeg, or .png."
+            log_entries.append(status)
+            logging.error(status)
+            return status
+
         log_entries.append(status)
         logging.info(status)
         return status
     except Exception as e:
-        video_loaded = False
-        status = f"Error loading video: {str(e)}"
+        media_loaded = False
+        status = f"Error loading media: {str(e)}"
         log_entries.append(status)
         logging.error(status)
         return status
@@ -137,11 +157,11 @@ def set_active_service(
     return None, "No Service Category Enabled"
 
 def generate_line_chart() -> Optional[str]:
-    if not crack_counts:
+    if not detected_counts:
         return None
     fig, ax = plt.subplots(figsize=(4, 2))
-    ax.plot(crack_counts[-50:], marker='o', color='#4682B4')
-    ax.set_title("Cracks/Holes Over Time")
+    ax.plot(detected_counts[-50:], marker='o', color='#4682B4')
+    ax.set_title("Detections Over Time")
     ax.set_xlabel("Frame")
     ax.set_ylabel("Count")
     ax.grid(True)
@@ -161,17 +181,18 @@ def monitor_feed() -> Tuple[
     Optional[str]
 ]:
     global paused, frame_count, last_frame, last_metrics, last_timestamp
-    global gps_coordinates, last_detected_cracks, last_detected_holes, video_loaded
+    global gps_coordinates, detected_plants, detected_issues, media_loaded
+    global is_video, static_image
 
-    if not video_loaded:
-        log_entries.append("Cannot start streaming: Video not loaded successfully.")
-        logging.error("Video not loaded successfully.")
+    if not media_loaded:
+        log_entries.append("Cannot start processing: Media not loaded successfully.")
+        logging.error("Media not loaded successfully.")
         return (
             None,
-            json.dumps({"error": "Video not loaded. Please upload a video file."}, indent=2),
+            json.dumps({"error": "Media not loaded. Please upload a video or image file."}, indent=2),
             "\n".join(log_entries[-10:]),
-            last_detected_cracks,
-            last_detected_holes,
+            detected_plants,
+            detected_issues,
             None,
             None
         )
@@ -181,9 +202,12 @@ def monitor_feed() -> Tuple[
         metrics = last_metrics.copy()
     else:
         try:
-            frame = get_next_video_frame()
-            if frame is None:
-                raise RuntimeError("Failed to retrieve frame from video.")
+            if is_video:
+                frame = get_next_video_frame()
+                if frame is None:
+                    raise RuntimeError("Failed to retrieve frame from video.")
+            else:
+                frame = static_image.copy()
         except RuntimeError as e:
             log_entries.append(f"Error: {str(e)}")
             logging.error(f"Frame retrieval error: {str(e)}")
@@ -191,8 +215,8 @@ def monitor_feed() -> Tuple[
                 None,
                 json.dumps(last_metrics, indent=2),
                 "\n".join(log_entries[-10:]),
-                last_detected_cracks,
-                last_detected_holes,
+                detected_plants,
+                detected_issues,
                 None,
                 None
             )
@@ -240,8 +264,28 @@ def monitor_feed() -> Tuple[
                 all_detected_items.extend(thermal_dets)
                 thermal_flag = bool(thermal_dets)
 
-            # Overlay detections
-            frame = overlay_boxes(frame, all_detected_items)
+            # Overlay detections with distinct colors
+            for item in all_detected_items:
+                box = item.get("box", [])
+                if not box:
+                    continue
+                x_min, y_min, x_max, y_max = box
+                label = item.get("label", "")
+                dtype = item.get("type", "")
+                # Assign colors based on detection type
+                if dtype == "plant":
+                    color = (0, 255, 0)  # Green for plants
+                elif dtype == "crack":
+                    color = (255, 0, 0)  # Red for cracks
+                elif dtype == "hole":
+                    color = (0, 0, 255)  # Blue for holes
+                elif dtype == "missing_patch":
+                    color = (255, 165, 0)  # Orange for missing patches
+                else:
+                    color = (255, 255, 0)  # Yellow for others
+                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
+                cv2.putText(frame, label, (x_min, y_min - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
             # Save temporary image
             cv2.imwrite(TEMP_IMAGE_PATH, frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
@@ -265,14 +309,15 @@ def monitor_feed() -> Tuple[
                 captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count}.jpg")
                 cv2.imwrite(captured_frame_path, frame)
                 for item in all_detected_items:
-                    if item.get("type") == "crack":
-                        last_detected_cracks.append(captured_frame_path)
-                        if len(last_detected_cracks) > 100:
-                            last_detected_cracks.pop(0)
-                    elif item.get("type") == "hole":
-                        last_detected_holes.append(captured_frame_path)
-                        if len(last_detected_holes) > 100:
-                            last_detected_holes.pop(0)
+                    dtype = item.get("type", "")
+                    if dtype == "plant":
+                        detected_plants.append(captured_frame_path)
+                        if len(detected_plants) > 100:
+                            detected_plants.pop(0)
+                    elif dtype in ["crack", "hole", "missing_patch"]:
+                        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)}")
@@ -309,29 +354,25 @@ def monitor_feed() -> Tuple[
         last_frame = frame.copy()
         last_metrics = metrics
 
-        # Track cracks/holes for metrics
+        # Track detections for metrics
+        plant_detected = len([item for item in all_detected_items if item.get("type") == "plant"])
         crack_detected = len([item for item in all_detected_items if item.get("type") == "crack"])
         hole_detected = len([item for item in all_detected_items if item.get("type") == "hole"])
-        if active_service in ["operations_maintenance", "road_safety"]:
-            crack_severity_all.extend([
-                item["severity"]
-                for item in all_detected_items
-                if item.get("type") in ["crack", "hole"] and "severity" in item
-            ])
-        crack_counts.append(crack_detected + hole_detected)
+        missing_detected = len([item for item in all_detected_items if item.get("type") == "missing_patch"])
+        detected_counts.append(plant_detected + crack_detected + hole_detected + missing_detected)
 
         # Log frame processing details
-        log_message = f"{last_timestamp} - Frame {frame_count} - Cracks: {crack_detected} - Holes: {hole_detected} - GPS: {gps_coord}"
+        log_message = (f"{last_timestamp} - Frame {frame_count} - "
+                       f"Plants: {plant_detected} - Cracks: {crack_detected} - "
+                       f"Holes: {hole_detected} - Missing Patches: {missing_detected} - GPS: {gps_coord}")
         log_entries.append(log_message)
         logging.info(log_message)
 
-        # Limit the size of logs and crack data
+        # Limit the size of logs and detection data
         if len(log_entries) > 100:
             log_entries.pop(0)
-        if len(crack_counts) > 500:
-            crack_counts.pop(0)
-        if len(crack_severity_all) > 500:
-            crack_severity_all.pop(0)
+        if len(detected_counts) > 500:
+            detected_counts.pop(0)
 
     # Resize frame and add metadata for display
     frame = cv2.resize(last_frame, (640, 480))
@@ -339,14 +380,15 @@ def monitor_feed() -> Tuple[
     cv2.putText(frame, f"{last_timestamp}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
 
     # Generate map
-    map_path = generate_map(gps_coordinates[-5:], [item for item in last_metrics.get("items", []) if item.get("type") in ["crack", "hole"]])
+    map_items = [item for item in last_metrics.get("items", []) if item.get("type") in ["crack", "hole", "missing_patch"]]
+    map_path = generate_map(gps_coordinates[-5:], map_items)
 
     return (
         frame[:, :, ::-1],  # Convert BGR to RGB for Gradio
         json.dumps(last_metrics, indent=2),
         "\n".join(log_entries[-10:]),
-        last_detected_cracks,
-        last_detected_holes,
+        detected_plants,
+        detected_issues,
         generate_line_chart(),
         map_path
     )
@@ -356,18 +398,18 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"))
     gr.Markdown(
         """
         # 🛡️ NHAI Drone Road Inspection Dashboard
-        Monitor highway conditions in real-time using drone footage. Select a service category to analyze specific aspects of the road.
+        Monitor highway conditions in real-time using drone footage or static images. Select a service category to analyze specific aspects of the road or plantation.
         """
     )
 
     with gr.Row():
         with gr.Column(scale=3):
-            video_input = gr.File(label="Upload Video File (e.g., sample.mp4)", file_types=[".mp4", ".avi"])
-            load_button = gr.Button("Load Video", variant="primary")
+            media_input = gr.File(label="Upload Media File (e.g., sample.mp4, image.jpg)", file_types=[".mp4", ".avi", ".jpg", ".jpeg", ".png"])
+            load_button = gr.Button("Load Media", variant="primary")
         with gr.Column(scale=1):
-            video_status = gr.Textbox(
-                label="Video Load Status",
-                value="Please upload a video file or ensure 'sample.mp4' exists in the root directory.",
+            media_status = gr.Textbox(
+                label="Media Load Status",
+                value="Please upload a video/image file or ensure 'sample.mp4' exists in the root directory.",
                 interactive=False
             )
 
@@ -385,29 +427,29 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"))
             pl_toggle = gr.Checkbox(label="Enable Plantation Services", value=False)
             pl_status = gr.Textbox(label="Plantation Status", value="Disabled", interactive=False)
 
-    status_text = gr.Markdown("**Status:** 🟢 Ready (Upload a video to start)")
+    status_text = gr.Markdown("**Status:** 🟢 Ready (Upload a media file to start)")
 
     with gr.Row():
         with gr.Column(scale=3):
-            video_output = gr.Image(label="Live Drone Feed", width=640, height=480, elem_id="live-feed")
+            media_output = gr.Image(label="Live Feed", width=640, height=480, elem_id="live-feed")
         with gr.Column(scale=1):
             metrics_output = gr.Textbox(
                 label="Detection Metrics",
                 lines=10,
                 interactive=False,
-                placeholder="Detection metrics, crack/hole counts will appear here."
+                placeholder="Detection metrics, counts will appear here."
             )
 
     with gr.Row():
         with gr.Column(scale=2):
             logs_output = gr.Textbox(label="Live Logs", lines=8, interactive=False)
         with gr.Column(scale=1):
-            crack_images = gr.Gallery(label="Detected Cracks (Last 100+)", columns=4, rows=13, height="auto")
-            hole_images = gr.Gallery(label="Detected Holes (Last 100+)", columns=4, rows=13, height="auto")
+            plant_images = gr.Gallery(label="Detected Plants (Last 100+)", columns=4, rows=13, height="auto")
+            issue_images = gr.Gallery(label="Detected Issues (Last 100+)", columns=4, rows=13, height="auto")
 
     with gr.Row():
-        chart_output = gr.Image(label="Crack/Hole Trend")
-        map_output = gr.Image(label="Crack/Hole Locations Map")
+        chart_output = gr.Image(label="Detection Trend")
+        map_output = gr.Image(label="Issue Locations Map")
 
     with gr.Row():
         pause_btn = gr.Button("⏸️ Pause", variant="secondary")
@@ -443,12 +485,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"))
         global frame_rate
         frame_rate = val
 
-    video_status.value = initialize_video()
+    media_status.value = initialize_media()
 
     load_button.click(
-        initialize_video,
-        inputs=[video_input],
-        outputs=[video_status]
+        initialize_media,
+        inputs=[media_input],
+        outputs=[media_status]
     )
 
     def update_toggles(uc_val: bool, om_val: bool, rs_val: bool, pl_val: bool) -> Tuple[str, str, str, str, str]:
@@ -474,17 +516,20 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"))
 
     def streaming_loop():
         while True:
-            if not video_loaded:
-                yield None, json.dumps({"error": "Video not loaded. Please upload a video file."}, indent=2), "\n".join(log_entries[-10:]), last_detected_cracks, last_detected_holes, None, None
+            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_plants, detected_issues, None, None
             else:
-                frame, metrics, logs, cracks, holes, chart, map_path = monitor_feed()
+                frame, metrics, logs, plants, issues, chart, map_path = monitor_feed()
                 if frame is None:
-                    yield None, metrics, logs, cracks, holes, chart, map_path
+                    yield None, metrics, logs, plants, issues, chart, map_path
                 else:
-                    yield frame, metrics, logs, cracks, holes, chart, map_path
+                    yield frame, metrics, logs, plants, issues, chart, map_path
+                if not is_video:
+                    # For static images, yield once and pause
+                    break
             time.sleep(frame_rate)
 
-    app.load(streaming_loop, outputs=[video_output, metrics_output, logs_output, crack_images, hole_images, chart_output, map_output])
+    app.load(streaming_loop, outputs=[media_output, metrics_output, logs_output, plant_images, issue_images, chart_output, map_output])
 
 if __name__ == "__main__":
     app.launch(share=False)