# camera_capture.py
import os
from pypylon import pylon
import cv2
from datetime import datetime


def search_get_device():
    """搜索并返回 Basler GigE 相机对象"""
    tl_factory = pylon.TlFactory.GetInstance()
    for dev_info in tl_factory.EnumerateDevices():
        if dev_info.GetDeviceClass() == 'BaslerGigE':
            print(f"Found Basler GigE Camera: {dev_info.GetModelName()}  IP: {dev_info.GetIpAddress()}")
            return pylon.InstantCamera(tl_factory.CreateDevice(dev_info))
    raise EnvironmentError("没有找到 Basler GigE 相机")


def grab_image(save=True):
    """
    采集一张图像并返回 numpy ndarray（BGR）
    save=True 时自动保存 PNG 文件到 camera/ 文件夹
    """
    # 创建保存目录
    save_dir = "camera"
    os.makedirs(save_dir, exist_ok=True)

    cam = search_get_device()
    cam.Open()

    # 设置图像转换器（BGR8 → OpenCV格式）
    converter = pylon.ImageFormatConverter()
    converter.OutputPixelFormat = pylon.PixelType_BGR8packed
    converter.OutputBitAlignment = pylon.OutputBitAlignment_MsbAligned

    cam.StartGrabbing(pylon.GrabStrategy_LatestImageOnly)

    grabResult = cam.RetrieveResult(5000, pylon.TimeoutHandling_ThrowException)

    if grabResult.GrabSucceeded():
        image = converter.Convert(grabResult)
        img = image.GetArray()     # OpenCV 可处理的 ndarray
        print("成功采集图像:", img.shape)

        # 保存 PNG
        if save:
            filename = os.path.join(save_dir, f"Image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png")
            cv2.imwrite(filename, img)
            print("已保存:", filename)

    else:
        print("采集失败:", grabResult.ErrorCode, grabResult.ErrorDescription)
        img = None

    grabResult.Release()
    cam.Close()
    return img


if __name__ == "__main__":
    img = grab_image(save=True)
    if img is not None:
        print("图像类型:", type(img))
        print("图像尺寸:", img.shape)