AVM360/1.py

import cv2
import threading
import sys
import os
import argparse
import numpy as np
from surround_view import FisheyeCameraModel, BirdView
import surround_view.param_settings as settings

class MultiCameraBirdView:
    def __init__(self):
        self.running = True
        self.names = settings.camera_names
        self.yamls = [os.path.join(os.getcwd(), "yaml", name + ".yaml") for name in self.names]
        self.camera_models = [FisheyeCameraModel(camera_file, camera_name) for camera_file, camera_name in zip(self.yamls, self.names)]
        
        # 初始化摄像头
        self.caps = []
        self.which_cameras = {
            "front": 0,
            "left": 1,
            "back": 2,
            "right":3
        }
        
        # 初始化摄像头捕获
        for name in self.names:
            cap = cv2.VideoCapture(self.which_cameras[name], cv2.CAP_V4L2)
            cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV"))
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
            cap.set(cv2.CAP_PROP_BUFFERSIZE, 3)
            if not cap.isOpened():
                print(f"[ERROR] Cannot open {name} camera", file=sys.stderr)
                self.running = False
                return
            self.caps.append(cap)
        
        # 初始化鸟瞰图
        self.birdview = BirdView()
        
        # 预先加载一次静态图像来初始化权重矩阵
        self._initialize_weights()
    
    def _initialize_weights(self):
        """使用静态图像初始化权重矩阵"""
        try:
            # 加载静态图像用于初始化权重
            images = [os.path.join(os.getcwd(), "images", name + ".png") for name in self.names]
            static_frames = []
            
            for image_file, camera_model in zip(images, self.camera_models):
                img = cv2.imread(image_file)
                if img is None:
                    print(f"[ERROR] Cannot load static image: {image_file}")
                    continue
                img = camera_model.undistort(img)
                img = camera_model.project(img)
                img = camera_model.flip(img)
                static_frames.append(img)
            
            if len(static_frames) == 4:
                # 初始化权重和掩码矩阵
                Gmat, Mmat = self.birdview.get_weights_and_masks(static_frames)
                print("[INFO] Weights and masks initialized successfully")
            else:
                print("[WARNING] Could not load all static images for weight initialization")
        except Exception as e:
            print(f"[ERROR] Failed to initialize weights: {e}", file=sys.stderr)
    
    def process_frame(self, img, camera_model):
        """处理单帧图像：去畸变 -> 投影 -> 翻转"""
        img = camera_model.undistort(img)
        img = camera_model.project(img)
        img = camera_model.flip(img)
        return img
    
    def run(self):
        # 获取屏幕分辨率（假设为 1920x1080，如果不同可以手动指定）
        # 在 Linux 下全屏显示通常需要先创建一个窗口
        # cv2.namedWindow("SurroundView_System", cv2.WND_PROP_FULLSCREEN)
        # cv2.setWindowProperty("SurroundView_System", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

        cv2.namedWindow("SurroundView_System", cv2.WINDOW_NORMAL)

        # 2. 指定窗口大小 (宽度, 高度)
        cv2.resizeWindow("SurroundView_System", 1024, 600)

        SCREEN_W = 1024
        SCREEN_H = 600
        
        # 计算布局比例
        # 鸟瞰图占 1/3 宽度: 1920 / 3 = 640
        # 原始画面占 2/3 宽度: 1920 * 2 / 3 = 1280
        BIRD_W = SCREEN_W // 3
        FRONT_W = SCREEN_W - BIRD_W
        
        while self.running:
            frames = []
            raw_front_frame = None
            
            for i, (cap, camera_model) in enumerate(zip(self.caps, self.camera_models)):
                ret, frame = cap.read()
                if not ret or frame is None:
                    continue
                
                # 保存一份原始的前向画面用于监视 (假设 front 是第一个摄像头)
                if self.names[i] == "front":
                    raw_front_frame = frame.copy()
                
                processed_frame = self.process_frame(frame, camera_model)
                frames.append(processed_frame)
            
            if len(frames) == 4 and raw_front_frame is not None:
                try:
                    # 1. 处理鸟瞰图
                    self.birdview.update_frames(frames)
                    self.birdview.stitch_all_parts()
                    self.birdview.make_white_balance()
                    self.birdview.copy_car_image()
                    
                    # 2. 缩放鸟瞰图 (保持比例，填满左侧 1/3)
                    # 假设原始 birdview.image 比较长，我们将其 resize 到 BIRD_W x SCREEN_H
                    bird_part = cv2.resize(self.birdview.image, (BIRD_W, SCREEN_H))
                    
                    # 3. 处理前向原始画面 (缩放至右侧 2/3)
                    # raw_front_frame 是 1280x720, 缩放到 FRONT_W x SCREEN_H
                    front_part = cv2.resize(raw_front_frame, (FRONT_W, SCREEN_H))
                    
                    # 4. 水平拼接 (左鸟瞰 + 右前视)
                    display_all = np.hstack((bird_part, front_part))
                    
                    # 5. 全屏显示
                    cv2.imshow("SurroundView_System", display_all)
                    
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        self.running = False
                        break
                        
                except Exception as e:
                    print(f"[ERROR] Processing error: {e}")
                    continue
        
        for cap in self.caps:
            cap.release()
        cv2.destroyAllWindows()



def main():
    parser = argparse.ArgumentParser(description="Real-time 4-Camera Bird's Eye View")
    parser.add_argument("--mode", type=str, required=True, 
                        choices=["realtime", "static"],
                        help="Mode: 'realtime' for live video, 'static' for static images")
    args = parser.parse_args()
    
    if args.mode == "static":
        # 原有的静态图像处理代码
        names = settings.camera_names
        images = [os.path.join(os.getcwd(), "images", name + ".png") for name in names]
        yamls = [os.path.join(os.getcwd(), "yaml", name + ".yaml") for name in names]
        camera_models = [FisheyeCameraModel(camera_file, camera_name) for camera_file, camera_name in zip(yamls, names)]

        projected = []
        for image_file, camera in zip(images, camera_models):
            img = cv2.imread(image_file)
            img = camera.undistort(img)
            img = camera.project(img)
            img = camera.flip(img)
            projected.append(img)

        birdview = BirdView()
        Gmat, Mmat = birdview.get_weights_and_masks(projected)  # 初始化权重矩阵
        birdview.update_frames(projected)
        # birdview.get_weights_and_masks.stitch_all_parts()
        # birdview.make_white_balance()
        birdview.copy_car_image()
        img_small = cv2.resize(birdview.image, (600, 1024))
        cv2.imshow("BirdView Result", img_small)
        
        while True:
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                cv2.destroyAllWindows()
                return -1
    
    elif args.mode == "realtime":
        # 实时视频处理
        print("Starting real-time 4-camera bird's eye view...")
        print("Press 'q' to quit")
        multi_cam = MultiCameraBirdView()
        if multi_cam.running:
            multi_cam.run()
        else:
            print("[ERROR] Failed to initialize cameras")


if __name__ == "__main__":
    main()