upload inference and app python scripts

app.py

@@ -0,0 +1,124 @@
+# app.py
+import time
+import cv2
+import numpy as np
+import gradio as gr
+
+from inference import CoralSegModel
+
+model = CoralSegModel()
+
+############################
+# 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
+
+    cap = cv2.VideoCapture(rtsp_or_http_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:
+                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
+            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.
+    """
+    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:
+                idx += 1
+                continue
+            processed = model.predict_overlay(frame)
+            yield processed
+            idx += 1
+            # tiny sleep to reduce CPU spikes; tune as needed
+            time.sleep(0.001)
+    finally:
+        cap.release()
+
+############################
+# 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**."
+    )
+
+    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.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")
+
+        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)
+
+if __name__ == "__main__":
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)

inference.py

@@ -0,0 +1,44 @@
+# inference.py
+import torch
+import numpy as np
+from PIL import Image
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+
+# Load model from HF (swap this with your own if you want)
+HF_MODEL_ID = "EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024"
+
+class CoralSegModel:
+    def __init__(self, device=None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.processor = SegformerImageProcessor.from_pretrained(HF_MODEL_ID)
+        self.model = SegformerForSemanticSegmentation.from_pretrained(HF_MODEL_ID).to(self.device)
+        self.model.eval()
+
+        # Build a simple color palette for masks (fallback if none provided)
+        # 0..N-1 colors - here random-ish but stable
+        num_classes = self.model.config.id2label and len(self.model.config.id2label) or 40
+        rng = np.random.RandomState(0)
+        self.palette = (rng.randint(0, 255, size=(num_classes, 3))).astype(np.uint8)
+
+    @torch.inference_mode()
+    def predict_overlay(self, frame_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
+        """
+        frame_bgr: np.ndarray HxWx3 in BGR (as read by OpenCV)
+        returns: np.ndarray HxWx3 in BGR (overlay)
+        """
+        # Convert BGR -> RGB PIL
+        rgb = frame_bgr[:, :, ::-1]
+        pil = Image.fromarray(rgb)
+
+        inputs = self.processor(images=pil, return_tensors="pt").to(self.device)
+        outputs = self.model(**inputs)
+        logits = outputs.logits  # [B, C, h, w]
+        upsampled = torch.nn.functional.interpolate(
+            logits, size=pil.size[::-1], mode="bilinear", align_corners=False
+        )
+        pred = upsampled.argmax(dim=1)[0].detach().cpu().numpy().astype(np.uint8)  # HxW
+
+        color_mask = self.palette[pred]  # HxWx3 (RGB)
+        overlay_rgb = (rgb * (1 - alpha) + color_mask * alpha).astype(np.uint8)
+        overlay_bgr = overlay_rgb[:, :, ::-1]
+        return overlay_bgr