MARLIN: An Edge-AI AR fusion for Marine Ecosystem Monitoring

"""
MARLIN Edge Node (sample code)
- Captures frames from a USB/CSI camera
- Records short audio windows from a USB hydrophone (via PyAudio)
- Runs placeholder inference (replace with your models)
- Streams results over a local WebSocket and also saves to disk
Tested on Raspberry Pi 4/5 (or Jetson Nano) with Python 3.10+
Dependencies: opencv-python, numpy, sounddevice, websockets, onnxruntime (optional)
"""

import asyncio
import base64
import json
import time
from collections import deque
from dataclasses import dataclass
from typing import Deque, Dict, Any

import cv2
import numpy as np
import sounddevice as sd
import websockets

# ---- Config ----
CAMERA_INDEX = 0              # 0 for default camera
AUDIO_SAMPLE_RATE = 48000     # hydrophone / USB audio
AUDIO_WINDOW_SEC = 2.0        # seconds per analysis window
SERVER_WS = "ws://127.0.0.1:8765"  # local WebSocket server (see marlin_ws_server.py)
SEND_PREVIEW_EVERY = 5        # send a small preview frame every N frames
PREVIEW_WIDTH = 320

# Replace with real model calls
def analyze_frame(frame: np.ndarray) -> Dict[str, Any]:
    """Dummy coral/fish/debris detector. Replace with your model inference."""
    h, w = frame.shape[:2]
    # toy heuristic: estimate "fish density" by counting edges
    edges = cv2.Canny(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), 60, 120)
    density = float(np.count_nonzero(edges)) / (w * h)
    result = {
        "coral_health": max(0.0, 1.0 - 2.5 * density),  # fake score
        "fish_index": 5.0 * density,                    # fake index
        "debris_prob": min(1.0, 3.0 * density),         # fake probability
    }
    return result

def analyze_audio(audio_block: np.ndarray, sr: int) -> Dict[str, Any]:
    """Dummy whale/ship noise index using simple spectral energy bands."""
    # Compute magnitude spectrum
    fft = np.fft.rfft(audio_block * np.hanning(len(audio_block)))
    mag = np.abs(fft)
    freqs = np.fft.rfftfreq(len(audio_block), d=1.0/sr)

    def band_energy(lo, hi):
        idx = (freqs >= lo) & (freqs < hi)
        return float(np.mean(mag[idx]) if np.any(idx) else 0.0)

    whale_band = band_energy(15, 80)     # rough baleen whale band (very approximate)
    ship_band  = band_energy(80, 300)    # rough engine/prop noise band (very approximate)
    ambient    = band_energy(300, 2000)

    # Normalize
    total = whale_band + ship_band + ambient + 1e-9
    return {
        "whale_index": whale_band / total,
        "ship_noise_index": ship_band / total,
        "ambient_index": ambient / total
    }

@dataclass
class AudioRing:
    maxlen: int
    sr: int
    buf: Deque[float]

    def __init__(self, seconds: float, sr: int):
        self.sr = sr
        self.maxlen = int(seconds * sr)
        self.buf = deque(maxlen=self.maxlen)

    def extend(self, samples: np.ndarray):
        self.buf.extend(samples.tolist())

    def window(self) -> np.ndarray:
        if len(self.buf) < self.maxlen // 2:
            return None
        arr = np.array(self.buf, dtype=np.float32)
        return arr[-self.maxlen:]

async def run():
    # Camera
    cap = cv2.VideoCapture(CAMERA_INDEX)
    if not cap.isOpened():
        raise RuntimeError("Camera not found")

    # Audio
    ring = AudioRing(seconds=AUDIO_WINDOW_SEC, sr=AUDIO_SAMPLE_RATE)
    sd.default.samplerate = AUDIO_SAMPLE_RATE
    sd.default.channels = 1

    def audio_cb(indata, frames, time_info, status):
        if status:
            print("Audio status:", status)
        ring.extend(indata[:, 0])

    stream = sd.InputStream(callback=audio_cb)
    stream.start()

    # WebSocket
    ws = await websockets.connect(SERVER_WS)
    print("Connected to:", SERVER_WS)

    frame_count = 0
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                await asyncio.sleep(0.01)
                continue

            vision = analyze_frame(frame)
            audio_window = ring.window()
            audio = analyze_audio(audio_window, AUDIO_SAMPLE_RATE) if audio_window is not None else None

            payload = {
                "ts": time.time(),
                "vision": vision,
                "audio": audio,
            }

            # Occasionally attach a preview image (base64)
            if frame_count % SEND_PREVIEW_EVERY == 0:
                preview = cv2.resize(frame, (PREVIEW_WIDTH, int(frame.shape[0] * PREVIEW_WIDTH / frame.shape[1])))
                _, jpg = cv2.imencode('.jpg', preview, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
                payload["preview_jpeg_b64"] = base64.b64encode(jpg.tobytes()).decode('ascii')

            await ws.send(json.dumps(payload))
            frame_count += 1
    finally:
        stream.stop()
        cap.release()
        await ws.close()

if __name__ == "__main__":
    asyncio.run(run())