simulation/src/vision/camera_processor.py

#!/usr/bin/env python3
"""
Camera Processor - Handles Gazebo camera feeds for GPS-denied navigation.
Subscribes to gz-transport topics, processes frames, and displays via OpenCV.
Downstream modules (ArUco detector, optical flow) register as callbacks.
"""

import sys
import signal
import time
import numpy as np
import cv2
from gz.transport13 import Node
from gz.msgs10.image_pb2 import Image

# PixelFormatType enum values (from gz.msgs10 Image proto)
PF_RGB_INT8 = 3
PF_BGR_INT8 = 8
PF_R_FLOAT32 = 13


class CameraProcessor:
    def __init__(self, topics=None, show_gui=True):
        self.node = Node()
        self.show_gui = show_gui
        self.frames = {}
        self.callbacks = {}
        self.running = True

        if topics is None:
            topics = {
                "downward": "/uav/camera/downward",
                "gimbal": "/world/uav_ugv_search/model/iris_with_gimbal/model/gimbal/link/pitch_link/sensor/camera/image",
            }

        self.topics = topics

        for name, topic in self.topics.items():
            self.callbacks[name] = []
            ok = self.node.subscribe(Image, topic, self._make_gz_callback(name))
            if ok:
                print(f"[CAM] Subscribed: {name} -> {topic}")
            else:
                print(f"[CAM] Failed: {topic}")

        signal.signal(signal.SIGINT, lambda s, f: self.stop())

    def _make_gz_callback(self, name):
        def cb(msg):
            fmt = msg.pixel_format_type
            if fmt == PF_RGB_INT8:
                arr = np.frombuffer(msg.data, dtype=np.uint8).reshape(
                    msg.height, msg.width, 3
                )
                bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
            elif fmt == PF_BGR_INT8:
                arr = np.frombuffer(msg.data, dtype=np.uint8).reshape(
                    msg.height, msg.width, 3
                )
                bgr = arr
            elif fmt == PF_R_FLOAT32:
                arr = np.frombuffer(msg.data, dtype=np.float32).reshape(
                    msg.height, msg.width
                )
                normalized = cv2.normalize(arr, None, 0, 255, cv2.NORM_MINMAX)
                bgr = cv2.cvtColor(normalized.astype(np.uint8), cv2.COLOR_GRAY2BGR)
            else:
                return

            processed = self.process_image(bgr)
            self.frames[name] = processed

            for fn in self.callbacks.get(name, []):
                fn(name, processed)

        return cb

    def process_image(self, image):
        processed = image.copy()
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
        lab[:, :, 0] = clahe.apply(lab[:, :, 0])
        processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
        return processed

    def register_callback(self, camera_name, fn):
        if camera_name in self.callbacks:
            self.callbacks[camera_name].append(fn)
            print(f"[CAM] Registered callback for {camera_name}: {fn.__name__}")

    def stop(self):
        self.running = False

    def spin(self):
        print("[CAM] Running. Press 'q' to quit, 's' to screenshot.")
        while self.running:
            for name, frame in list(self.frames.items()):
                cv2.imshow(name, frame)

            key = cv2.waitKey(33) & 0xFF
            if key == ord("q"):
                break
            elif key == ord("s"):
                for name, frame in self.frames.items():
                    fname = f"{name}_{int(time.time())}.png"
                    cv2.imwrite(fname, frame)
                    print(f"[CAM] Saved {fname}")

        cv2.destroyAllWindows()
        print("[CAM] Stopped.")

    def spin_headless(self):
        print("[CAM] Running headless (no GUI).")
        while self.running:
            time.sleep(0.1)
        print("[CAM] Stopped.")

    def get_frame(self, camera_name):
        return self.frames.get(camera_name)


def main():
    cameras = "both"
    show_gui = True

    if len(sys.argv) > 1:
        cameras = sys.argv[1].lower()
    if "--headless" in sys.argv:
        show_gui = False

    all_topics = {
        "downward": "/uav/camera/downward",
        "gimbal": "/world/uav_ugv_search/model/iris_with_gimbal/model/gimbal/link/pitch_link/sensor/camera/image",
    }

    if cameras == "down":
        topics = {"downward": all_topics["downward"]}
    elif cameras == "gimbal":
        topics = {"gimbal": all_topics["gimbal"]}
    else:
        topics = all_topics

    proc = CameraProcessor(topics=topics, show_gui=show_gui)

    try:
        from vision.object_detector import ObjectDetector
        detector = ObjectDetector(aruco_dict_name="DICT_4X4_50")

        def detection_overlay(camera_name, frame):
            detections = detector.detect(frame)
            if detections:
                annotated = detector.annotate_frame(frame, detections)
                proc.frames[camera_name] = annotated

        for cam_name in topics:
            proc.register_callback(cam_name, detection_overlay)
    except Exception as e:
        print(f"[CAM] ArUco detection unavailable: {e}")

    if show_gui:
        proc.spin()
    else:
        proc.spin_headless()


if __name__ == "__main__":
    main()