Update app.py

app.py

@@ -1,124 +1,191 @@
-# app.py
-import time
+from PIL import Image
 import cv2
 import numpy as np
 import gradio as gr
 
-from inference import CoralSegModel
-
+from inference import CoralSegModel, id2label, label2color, create_segmentation_overlay
 model = CoralSegModel()
 
-############################
-# Helpers
-############################
+# ---- helpers ----
 def _safe_read(cap):
     ok, frame = cap.read()
-    if not ok or frame is None:
-        return None
-    return frame
-
-############################
-# 1) Remote stream (server-pull)
-############################
-def remote_stream(rtsp_or_http_url: str, skip_every_n=1):
-    """
-    Generator that yields processed frames for gr.Video streaming.
-    - rtsp_or_http_url: e.g., rtsp://..., http://mjpeg..., or a video file URL
-    """
-    if not rtsp_or_http_url:
-        yield None
+    return frame if ok and frame is not None else None
+
+def build_annotations(pred_map: np.ndarray, selected: list[str]) -> list[tuple[np.ndarray, str]]:
+    """Return [(mask,label), ...] where mask is 0/1 float HxW for AnnotatedImage."""
+    if pred_map is None or not selected:
+        return []
+    
+    # Create reverse mapping: label_name -> class_id
+    label2id = {label: int(id_str) for id_str, label in id2label.items()}
+    
+    anns = []
+    for label_name in selected:
+        if label_name not in label2id:
+            continue  # Skip unknown labels
+        
+        class_id = label2id[label_name]  # Convert label name to class ID
+        mask = (pred_map == class_id).astype(np.float32)
+        if mask.sum() > 0:
+            anns.append((mask, label_name))  # Use the label name for display
+    return anns
+    
+# ==============================
+# STREAMING EVENT FUNCTIONS
+# ==============================
+# IMPORTANT: make the event functions themselves generators.
+# Also: include the States as outputs so we can update them every frame.
+def remote_start(url: str, n: int, pred_state, base_state):
+    if not url:
         return
-
-    cap = cv2.VideoCapture(rtsp_or_http_url)
+    cap = cv2.VideoCapture(url)
     if not cap.isOpened():
-        yield None
         return
-
     idx = 0
     try:
         while True:
             frame = _safe_read(cap)
             if frame is None:
                 break
-
-            if skip_every_n > 1 and (idx % skip_every_n) != 0:
+            if n > 1 and (idx % n) != 0:
                 idx += 1
                 continue
-
-            processed = model.predict_overlay(frame)
-            # IMPORTANT: Gradio 5 streaming expects raw numpy frames (H, W, 3) BGR/RGB both supported for display
-            yield processed
+            pred_map, overlay_rgb, base_rgb = model.predict_map_and_overlay(frame)
+            # yield live image + updated States' *values*
+            yield overlay_rgb, pred_map, base_rgb
             idx += 1
-            # Lower CPU usage a bit (tune this)
-            # time.sleep(0.001)
     finally:
         cap.release()
 
-def uploaded_video_stream(video_file, skip_every_n=1):
-    """
-    Gradio passes the uploaded file path (string) for gr.Video.
-    We open it with OpenCV and yield processed frames to stream.
-    """
+def upload_start(video_file: str, n: int):
     if not video_file:
-        yield None
         return
-
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
-        yield None
         return
-
     idx = 0
     try:
         while True:
             ok, frame = cap.read()
             if not ok or frame is None:
                 break
-            if skip_every_n > 1 and (idx % skip_every_n) != 0:
+            if n > 1 and (idx % n) != 0:
                 idx += 1
                 continue
-            processed = model.predict_overlay(frame)
-            yield processed
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            pred_map, overlay_rgb, base_rgb = model.predict_map_and_overlay(frame)
+            yield overlay_rgb, pred_map, base_rgb
             idx += 1
-            # tiny sleep to reduce CPU spikes; tune as needed
-            # time.sleep(0.001)
     finally:
         cap.release()
 
-############################
+# ==============================
+# SNAPSHOT / TOGGLES (non-streaming)
+# ==============================
+# NOTE: When you pass gr.State as an input, you receive the *value*, not the wrapper.
+def make_snapshot(selected_labels, pred_map, base_rgb, alpha=0.25):
+    if pred_map is None or base_rgb is None:
+        return gr.update()
+    # rebuild overlay to match the live look
+    overlay = create_segmentation_overlay(pred_map, id2label, label2color, Image.fromarray(base_rgb), alpha=alpha)
+    ann = build_annotations(pred_map, selected_labels or [])
+    return (overlay, ann)  # (base_image, [(mask,label), ...])
+
+# ==============================
 # UI
-############################
+# ==============================
 with gr.Blocks(title="CoralScapes Streaming Segmentation") as demo:
     gr.Markdown("# CoralScapes Streaming Segmentation")
     gr.Markdown(
-        "Two modes: **Remote Stream** (paste RTSP/HTTP/MJPEG URL) or **Upload Video**."
+        "Left: **live stream** (fast). Right: **snapshot** with **hover labels** and **per-class toggles**."
     )
 
     with gr.Tab("Remote Stream (RTSP/HTTP)"):
-        url = gr.Textbox(
-            label="Stream URL (rtsp://..., http://...)", placeholder="rtsp://user:pass@ip:port/..."
-        )
-        skip = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame (perf tweak)")
-        out_image = gr.Image(label="Segmented Stream", streaming=True)  # Changed to Image
-        start_btn = gr.Button("Start")
-        stop_btn = gr.Button("Stop")
-
-        def _start(url_value, n):
-            return remote_stream(url_value, int(n))
-
-        start_btn.click(_start, inputs=[url, skip], outputs=out_image)
-        stop_btn.click(lambda: None, inputs=None, outputs=out_image)
+        with gr.Row():
+            with gr.Column(scale=2):
+                
+                # States start as None. We'll UPDATE them on every frame by returning them as outputs.
+                pred_state_remote = gr.State(None)  # holds last pred_map (HxW np.uint8)
+                base_state_remote = gr.State(None)  # holds last base_rgb (HxWx3 uint8)
+
+                live_remote = gr.Image(label="Live segmented stream")
+
+                start_btn = gr.Button("Start")
+
+                snap_btn_remote = gr.Button("📸 Snapshot (hover-able)")
+                hover_remote = gr.AnnotatedImage(label="Snapshot (hover to see label)")
+                
+
+            with gr.Column(scale=1):
+                url  = gr.Textbox(label="Stream URL", placeholder="rtsp://user:pass@ip:port/…")
+                skip = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame")
+
+                toggles_remote = gr.CheckboxGroup(
+                    choices=list(id2label.values()), value=list(id2label.values()),
+                    label="Toggle classes in snapshot",
+                )
+
+            start_btn.click(
+                    remote_start,
+                    inputs=[url, skip, pred_state_remote, base_state_remote],
+                    outputs=[live_remote, pred_state_remote, base_state_remote],
+                    queue=True,   # be explicit; required for generator streaming
+                )
+            
+            snap_btn_remote.click(
+                make_snapshot,
+                inputs=[toggles_remote, pred_state_remote, base_state_remote],
+                outputs=[hover_remote],
+            )
+            toggles_remote.change(
+                make_snapshot,
+                inputs=[toggles_remote, pred_state_remote, base_state_remote],
+                outputs=[hover_remote],
+            )
 
     with gr.Tab("Upload Video"):
-        gr.Markdown("Upload a video file; the server will stream segmented frames back in real time.")
-        vid_in = gr.Video(sources=["upload"], format="mp4", label="Input Video")
-        out_image = gr.Image(label="Segmented Output (streaming)", streaming=True)  # Changed to Image
-        start_btn2 = gr.Button("Process")
-        stop_btn2 = gr.Button("Stop")
+        with gr.Row():
+            # Left column (now contains toggles, snapshot button, and live output)
+            with gr.Column(scale=2):
+                # States remain in the same column as live_upload
+                pred_state_upload = gr.State(None)
+                base_state_upload = gr.State(None)
+                
+                live_upload = gr.Image(label="Live segmented output")
+                start_btn2 = gr.Button("Process")
+                
+                snap_btn_upload = gr.Button("📸 Snapshot (hover-able)")
+                hover_upload = gr.AnnotatedImage(label="Snapshot (hover to see label)")
+                
+            # Right column (now contains video input and slider)
+            with gr.Column(scale=1):
+                vid_in = gr.Video(sources=["upload"], format="mp4", label="Input Video")
+                skip2 = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame")
+
+                toggles_upload = gr.CheckboxGroup(
+                    choices=list(id2label.values()), value=list(id2label.values()),
+                    label="Toggle classes in snapshot",
+                )
+                
+        # Event handlers remain the same
+        start_btn2.click(
+            upload_start,
+            inputs=[vid_in, skip2],
+            outputs=[live_upload, pred_state_upload, base_state_upload],
+            queue=True,
+        )
 
-        skip2 = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame")
-        start_btn2.click(uploaded_video_stream, inputs=[vid_in, skip2], outputs=out_image)
-        stop_btn2.click(lambda: None, inputs=None, outputs=out_image)
+        snap_btn_upload.click(
+            make_snapshot,
+            inputs=[toggles_upload, pred_state_upload, base_state_upload],
+            outputs=[hover_upload],
+        )
+        
+        toggles_upload.change(
+            make_snapshot,
+            inputs=[toggles_upload, pred_state_upload, base_state_upload],
+            outputs=[hover_upload],
+        )
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(share=True)