LhatMjnk/CoralreefSegmentation
0
1from PIL import Image2import cv23import numpy as np4import gradio as gr5 6from inference import CoralSegModel, id2label, label2color, create_segmentation_overlay7model = CoralSegModel()8 9# ---- helpers ----10def _safe_read(cap):11 ok, frame = cap.read()12 return frame if ok and frame is not None else None13 14def build_annotations(pred_map: np.ndarray, selected: list[str]) -> list[tuple[np.ndarray, str]]:15 """Return [(mask,label), ...] where mask is 0/1 float HxW for AnnotatedImage."""16 if pred_map is None or not selected:17 return []18 19 # Create reverse mapping: label_name -> class_id20 label2id = {label: int(id_str) for id_str, label in id2label.items()}21 22 anns = []23 for label_name in selected:24 if label_name not in label2id:25 continue # Skip unknown labels26 27 class_id = label2id[label_name] # Convert label name to class ID28 mask = (pred_map == class_id).astype(np.float32)29 if mask.sum() > 0:30 anns.append((mask, label_name)) # Use the label name for display31 return anns32 33# ==============================34# STREAMING EVENT FUNCTIONS35# ==============================36# IMPORTANT: make the event functions themselves generators.37# Also: include the States as outputs so we can update them every frame.38def remote_start(url: str, n: int, pred_state, base_state):39 if not url:40 return41 cap = cv2.VideoCapture(url)42 if not cap.isOpened():43 return44 idx = 045 try:46 while True:47 frame = _safe_read(cap)48 if frame is None:49 break50 if n > 1 and (idx % n) != 0:51 idx += 152 continue53 pred_map, overlay_rgb, base_rgb = model.predict_map_and_overlay(frame)54 # yield live image + updated States' *values*55 yield overlay_rgb, pred_map, base_rgb56 idx += 157 finally:58 cap.release()59 60def upload_start(video_file: str, n: int):61 if not video_file:62 return63 cap = cv2.VideoCapture(video_file)64 if not cap.isOpened():65 return66 idx = 067 try:68 while True:69 ok, frame = cap.read()70 if not ok or frame is None:71 break72 if n > 1 and (idx % n) != 0:73 idx += 174 continue75 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)76 pred_map, overlay_rgb, base_rgb = model.predict_map_and_overlay(frame)77 yield overlay_rgb, pred_map, base_rgb78 idx += 179 finally:80 cap.release()81 82# ==============================83# SNAPSHOT / TOGGLES (non-streaming)84# ==============================85# NOTE: When you pass gr.State as an input, you receive the *value*, not the wrapper.86def make_snapshot(selected_labels, pred_map, base_rgb, alpha=0.25):87 if pred_map is None or base_rgb is None:88 return gr.update()89 # rebuild overlay to match the live look90 overlay = create_segmentation_overlay(pred_map, id2label, label2color, Image.fromarray(base_rgb), alpha=alpha)91 ann = build_annotations(pred_map, selected_labels or [])92 return (overlay, ann) # (base_image, [(mask,label), ...])93 94# ==============================95# UI96# ==============================97with gr.Blocks(title="CoralScapes Streaming Segmentation") as demo:98 gr.Markdown("# CoralScapes Streaming Segmentation")99 gr.Markdown(100 "Left: **live stream** (fast). Right: **snapshot** with **hover labels** and **per-class toggles**."101 )102 103 with gr.Tab("Remote Stream (RTSP/HTTP)"):104 with gr.Row():105 with gr.Column(scale=2):106 107 # States start as None. We'll UPDATE them on every frame by returning them as outputs.108 pred_state_remote = gr.State(None) # holds last pred_map (HxW np.uint8)109 base_state_remote = gr.State(None) # holds last base_rgb (HxWx3 uint8)110 111 live_remote = gr.Image(label="Live segmented stream")112 113 start_btn = gr.Button("Start")114 115 snap_btn_remote = gr.Button("📸 Snapshot (hover-able)")116 hover_remote = gr.AnnotatedImage(label="Snapshot (hover to see label)")117 118 119 with gr.Column(scale=1):120 url = gr.Textbox(label="Stream URL", placeholder="rtsp://user:pass@ip:port/…")121 skip = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame")122 123 toggles_remote = gr.CheckboxGroup(124 choices=list(id2label.values()), value=list(id2label.values()),125 label="Toggle classes in snapshot",126 )127 128 start_btn.click(129 remote_start,130 inputs=[url, skip, pred_state_remote, base_state_remote],131 outputs=[live_remote, pred_state_remote, base_state_remote],132 queue=True, # be explicit; required for generator streaming133 )134 135 snap_btn_remote.click(136 make_snapshot,137 inputs=[toggles_remote, pred_state_remote, base_state_remote],138 outputs=[hover_remote],139 )140 toggles_remote.change(141 make_snapshot,142 inputs=[toggles_remote, pred_state_remote, base_state_remote],143 outputs=[hover_remote],144 )145 146 with gr.Tab("Upload Video"):147 with gr.Row():148 # Left column (now contains toggles, snapshot button, and live output)149 with gr.Column(scale=2):150 # States remain in the same column as live_upload151 pred_state_upload = gr.State(None)152 base_state_upload = gr.State(None)153 154 live_upload = gr.Image(label="Live segmented output")155 start_btn2 = gr.Button("Process")156 157 snap_btn_upload = gr.Button("📸 Snapshot (hover-able)")158 hover_upload = gr.AnnotatedImage(label="Snapshot (hover to see label)")159 160 # Right column (now contains video input and slider)161 with gr.Column(scale=1):162 vid_in = gr.Video(sources=["upload"], format="mp4", label="Input Video")163 skip2 = gr.Slider(1, 5, value=1, step=1, label="Process every Nth frame")164 165 toggles_upload = gr.CheckboxGroup(166 choices=list(id2label.values()), value=list(id2label.values()),167 label="Toggle classes in snapshot",168 )169 170 # Event handlers remain the same171 start_btn2.click(172 upload_start,173 inputs=[vid_in, skip2],174 outputs=[live_upload, pred_state_upload, base_state_upload],175 queue=True,176 )177 178 snap_btn_upload.click(179 make_snapshot,180 inputs=[toggles_upload, pred_state_upload, base_state_upload],181 outputs=[hover_upload],182 )183 184 toggles_upload.change(185 make_snapshot,186 inputs=[toggles_upload, pred_state_upload, base_state_upload],187 outputs=[hover_upload],188 )189 190if __name__ == "__main__":191 demo.queue().launch(share=True)192