smart-bow-infrared/Assets/LightGlue/Scripts/Roma/RomaForwardedImageViewer.cs

using System;
using UnityEngine;
using UnityEngine.SceneManagement;
using UnityEngine.UI;

namespace LightGlue.Unity.Roma
{
    /// <summary>
    /// 接收 Python 转发过来的 WiFi 图像流（FrameHeader 分片协议），并显示到 RawImage。
    /// Roma 新场景专用。
    /// </summary>
    public sealed class RomaForwardedImageViewer : MonoBehaviour
    {
        public static RomaForwardedImageViewer Instance { get; private set; }

        [Header("Receive (Python -> Unity image forward)")]
        public string bindIp = "0.0.0.0";

        [Tooltip("Python 转发到 Unity 的图像端口（与 Python --forward_port 对齐）")]
        public int forwardPort = 12366;

        [Tooltip("最大队列大小（只保留最新N帧）")]
        public int maxQueueSize = 2;

        [Tooltip("在途帧上限（防止网络抖动时占用过多内存）")]
        public int maxInFlightFrames = 8;

        [Tooltip("分片重组超时（ms）")]
        public int frameTimeoutMs = 500;

        [Header("UI")]
        public RawImage targetRawImage;

        [Tooltip("若 RawImage 未赋值，是否自动在场景中寻找第一个 RawImage")]
        public bool autoFindRawImage = true;

        [Tooltip("切场景后若 RawImage 被销毁，是否自动尝试重新绑定（FindObjectOfType<RawImage>）。")]
        public bool autoRebindRawImageOnSceneLoaded = true;

        [Tooltip("若始终找不到 RawImage，是否自动创建一个常驻的全屏 RawImage 用于显示。")]
        public bool autoCreatePersistentRawImage = true;

        [Tooltip("自动创建的常驻显示根节点名称")]
        public string persistentDisplayRootName = "RomaForwardedImageDisplay";

        [Header("Debug")]
        public bool logDecodeFailure = true;
        public float decodeFailureLogIntervalSeconds = 1.0f;

        [Header("Lifecycle")]
        [Tooltip("是否在 OnEnable 自动开始监听转发端口。若为 false，需要显式调用 StartReceiver()。")]
        public bool autoStartOnEnable = false;

        [Tooltip("是否让该组件所在 GameObject 常驻（DontDestroyOnLoad）。建议开启，避免退回场景后丢图像。")]
        public bool dontDestroyOnLoad = true;

        private RomaFrameHeaderJpegReceiver _receiver;
        private Texture2D _tex;
        private float _lastFailLogTime;

        private GameObject _persistentDisplayRoot;

        private void Awake()
        {
            if (Instance != null && Instance != this)
            {
                // 已有常驻实例在运行：销毁重复的，避免重复占用同一 UDP 端口
                Destroy(gameObject);
                return;
            }

            Instance = this;
            if (dontDestroyOnLoad)
                DontDestroyOnLoad(gameObject);
        }

        private void OnEnable()
        {
            SceneManager.sceneLoaded -= OnSceneLoaded;
            SceneManager.sceneLoaded += OnSceneLoaded;

            if (autoStartOnEnable)
            {
                StartReceiver();
            }
        }

        private void OnDisable()
        {
            SceneManager.sceneLoaded -= OnSceneLoaded;
            StopReceiver();
        }

        private void OnDestroy()
        {
            SceneManager.sceneLoaded -= OnSceneLoaded;
            if (Instance == this) Instance = null;
        }

        public void StartReceiver()
        {
            if (_receiver != null) return;

            EnsureRawImageBound();

            try
            {
                _receiver = new RomaFrameHeaderJpegReceiver(
                    bindIp: bindIp,
                    port: forwardPort,
                    maxQueueSize: maxQueueSize,
                    maxInFlightFrames: maxInFlightFrames,
                    frameTimeoutMs: frameTimeoutMs);
                _receiver.Start();
            }
            catch (Exception ex)
            {
                // 典型原因：场景中出现重复实例，或端口被其他组件占用
                Debug.LogError($"[RomaViewer] StartReceiver failed on {bindIp}:{forwardPort}: {ex.Message}");
                try { _receiver?.Stop(); } catch { /* ignore */ }
                _receiver = null;
                return;
            }

            if (_tex == null)
                _tex = new Texture2D(2, 2, TextureFormat.RGB24, false);
        }

        public void StopReceiver()
        {
            try { _receiver?.Stop(); } catch { /* ignore */ }
            _receiver = null;
        }

        private void Update()
        {
            if (_receiver == null) return;

            byte[] latest = null;
            while (_receiver.TryDequeueJpeg(out var jpeg))
                latest = jpeg;

            if (latest == null || latest.Length < 4) return;
            if (targetRawImage == null)
            {
                // 场景切换后 UI 被销毁：尽量自动恢复显示
                EnsureRawImageBound();
                if (targetRawImage == null) return;
            }

            // 最基本校验：JPEG SOI/EOI
            if (latest[0] != 0xFF || latest[1] != 0xD8 || latest[^2] != 0xFF || latest[^1] != 0xD9)
                return;

            bool ok = false;
            try
            {
                ok = _tex.LoadImage(latest, markNonReadable: false);
            }
            catch (Exception)
            {
                ok = false;
            }

            if (ok)
            {
                targetRawImage.texture = _tex;
            }
            else if (logDecodeFailure && (Time.time - _lastFailLogTime) >= decodeFailureLogIntervalSeconds)
            {
                _lastFailLogTime = Time.time;
                Debug.LogWarning($"[RomaViewer] Failed to decode JPEG (size={latest.Length} bytes) from {bindIp}:{forwardPort}");
            }
        }

        private void OnSceneLoaded(Scene scene, LoadSceneMode mode)
        {
            if (!autoRebindRawImageOnSceneLoaded) return;
            if (targetRawImage != null) return;
            EnsureRawImageBound();
        }

        private void EnsureRawImageBound()
        {
            if (targetRawImage != null) return;

            if (autoFindRawImage)
                targetRawImage = FindObjectOfType<RawImage>();

            if (targetRawImage != null) return;
            if (!autoCreatePersistentRawImage) return;

            if (_persistentDisplayRoot == null)
            {
                _persistentDisplayRoot = new GameObject(persistentDisplayRootName);
                DontDestroyOnLoad(_persistentDisplayRoot);

                var canvas = _persistentDisplayRoot.AddComponent<Canvas>();
                canvas.renderMode = RenderMode.ScreenSpaceOverlay;
                canvas.sortingOrder = short.MaxValue;

                _persistentDisplayRoot.AddComponent<CanvasScaler>();
                _persistentDisplayRoot.AddComponent<GraphicRaycaster>();

                var imgGo = new GameObject("RawImage");
                imgGo.transform.SetParent(_persistentDisplayRoot.transform, false);
                var raw = imgGo.AddComponent<RawImage>();

                var rt = raw.rectTransform;
                rt.anchorMin = Vector2.zero;
                rt.anchorMax = Vector2.one;
                rt.offsetMin = Vector2.zero;
                rt.offsetMax = Vector2.zero;

                raw.raycastTarget = false;
                targetRawImage = raw;
            }
            else
            {
                targetRawImage = _persistentDisplayRoot.GetComponentInChildren<RawImage>(true);
            }
        }
    }
}