from pathlib import Path

import cv2
from depthai import CameraBoardSocket
from depthai import ColorCameraProperties
from depthai import Device
from depthai import IMUSensor
from depthai import ImgFrame
from depthai import MonoCameraProperties
from depthai import Pipeline
from depthai import RawStereoDepthConfig
from depthai import StereoDepthProperties
from depthai import VideoEncoderProperties


HI_RES_STREAM_NAME = "hi_res"
COLOR_STREAM_NAME = "color"
CONTROL_STREAM_NAME = "control"
RIGHT_STREAM_NAME = "right"
DEPTH_COLOR_STREAM_NAME = "depth_color_aligned"
DEPTH_RIGHT_STREAM_NAME = "depth_right_aligned"
FEATURES_STREAM_NAME = "features"
DETECTIONS_STREAM_NAME = "detections"
IMU_STREAM_NAME = "imu"
HI_RES_IMAGE_SCRIPT_ISTREAM = "IN"
HI_RES_IMAGE_SCRIPT_OSTREAM = "OUT"


def create_hi_res_image_script(istream: str, ostream: str, modulo: int) -> str:
    return f"""
    i = 0
    while True:
        image = node.io["{istream}"].get()
        i += 1
        i %= {modulo}
        if i == 0:
            node.io["{ostream}"].send(image)
    """


pipeline = Pipeline()

color = pipeline.createColorCamera()
color.setBoardSocket(CameraBoardSocket.RGB)
color.setFps(5.0)
color.setResolution(ColorCameraProperties.SensorResolution.THE_4_K)
color.setVideoSize(2144, 2144)
color.setInterleaved(False)
color.initialControl.setSharpness(0)
color.initialControl.setLumaDenoise(0)
color.initialControl.setChromaDenoise(1)
color.initialControl.setAutoFocusLensRange(0, 255)
color.initialControl.setAutoExposureLimit(1000)

left = pipeline.createMonoCamera()
left.setBoardSocket(CameraBoardSocket.LEFT)
left.setFps(10.0)
left.setResolution(MonoCameraProperties.SensorResolution.THE_400_P)

right = pipeline.createMonoCamera()
right.setBoardSocket(CameraBoardSocket.RIGHT)
right.setFps(10.0)
right.setResolution(MonoCameraProperties.SensorResolution.THE_400_P)

depth = pipeline.createStereoDepth()
depth.setLeftRightCheck(True)
depth.initialConfig.setConfidenceThreshold(180)
depth.initialConfig.setMedianFilter(StereoDepthProperties.MedianFilter.MEDIAN_OFF)
depth.initialConfig.setSubpixel(True)
depth.initialConfig.setBilateralFilterSigma(0)
depth.initialConfig.setSubpixelFractionalBits(5)
depth.initialConfig.setDisparityShift(30)
depth.initialConfig.setDepthUnit(RawStereoDepthConfig.AlgorithmControl.DepthUnit.CUSTOM)
depth_config = depth.initialConfig.get()
depth_config.algorithmControl.customDepthUnitMultiplier = 50000
depth.initialConfig.set(depth_config)
depth.enableDistortionCorrection(True)

manip = pipeline.createImageManip()
manip.initialConfig.setResize(448, 448)
manip.initialConfig.setFrameType(ImgFrame.Type.BGR888p)

hi_res_image_script = pipeline.createScript()
hi_res_image_script.setScript(create_hi_res_image_script(HI_RES_IMAGE_SCRIPT_ISTREAM, HI_RES_IMAGE_SCRIPT_OSTREAM, 40))

video_encoder = pipeline.createVideoEncoder()
video_encoder.setProfile(VideoEncoderProperties.Profile.MJPEG)
video_encoder.setQuality(90)

features = pipeline.createFeatureTracker()
features.setHardwareResources(numShaves=2, numMemorySlices=2)

detection = pipeline.createYoloDetectionNetwork()
detection.setBlobPath(Path.home() / 'path' / 'to' / 'blob')
detection.setConfidenceThreshold(0.1)
detection.setNumClasses(2)
detection.setCoordinateSize(4)
detection.setIouThreshold(0.5)
detection.setNumInferenceThreads(2)
detection.input.setBlocking(False)
detection.input.setQueueSize(1)

imu = pipeline.createIMU()
imu.enableIMUSensor(IMUSensor.GYROSCOPE_RAW, 50)
imu.enableIMUSensor(IMUSensor.ACCELEROMETER_RAW, 50)
imu.setBatchReportThreshold(1)
imu.setMaxBatchReports(10)

control_in = pipeline.createXLinkIn()
hi_res_out = pipeline.createXLinkOut()
color_out = pipeline.createXLinkOut()
right_out = pipeline.createXLinkOut()
depth_out = pipeline.createXLinkOut()
features_out = pipeline.createXLinkOut()
detection_out = pipeline.createXLinkOut()
imu_out = pipeline.createXLinkOut()

control_in.setStreamName(CONTROL_STREAM_NAME)
hi_res_out.setStreamName(HI_RES_STREAM_NAME)
color_out.setStreamName(COLOR_STREAM_NAME)
right_out.setStreamName(RIGHT_STREAM_NAME)
depth_out.setStreamName(DEPTH_RIGHT_STREAM_NAME)
features_out.setStreamName(FEATURES_STREAM_NAME)
detection_out.setStreamName(DETECTIONS_STREAM_NAME)
imu_out.setStreamName(IMU_STREAM_NAME)

control_in.out.link(color.inputControl)
color.video.link(manip.inputImage)
color.video.link(hi_res_image_script.inputs[HI_RES_IMAGE_SCRIPT_ISTREAM])
manip.out.link(color_out.input)
manip.out.link(detection.input)
hi_res_image_script.outputs[HI_RES_IMAGE_SCRIPT_OSTREAM].link(video_encoder.input)
video_encoder.bitstream.link(hi_res_out.input)
left.out.link(depth.left)
right.out.link(depth.right)
depth.depth.link(depth_out.input)
depth.rectifiedRight.link(features.inputImage)
features.outputFeatures.link(features_out.input)
detection.out.link(detection_out.input)
depth.rectifiedRight.link(right_out.input)
imu.out.link(imu_out.input)

depth_out.input.setBlocking(False)
depth_out.input.setQueueSize(1)


#


with Device(pipeline) as device:
    color_queue = device.getOutputQueue(name=COLOR_STREAM_NAME, maxSize=4, blocking=False)
    while True:
        color_frame = color_queue.get().getCvFrame()
        print(str(device.readCalibration().getCameraIntrinsics(CameraBoardSocket.RGB)))
        print(str(device.readCalibration().getDistortionCoefficients(CameraBoardSocket.RGB)))
        cv2.imshow('color', color_frame)
        cv2.waitKey(100)

    WouterOddBot
    I can repro your issue and it seems to me that it only happens if I use the setAutoFocusLensRange. Removing that fixes the issue. Can you confirm so I can pass it to the FW team?

    Thanks,
    Jaka

      Hi @jakaskerl ,
      Yeah it looks like it. But I'm not entirely sure, as we actually started using setAutoFocusLensRange to see if it would help solving the focus problem…

      Apart from that, could you also explain a bit about my second question that the camera intrinsics don't change if the camera changes focus. I'd expect the camera intrinsics to change of the camera zooms in or out. Is that correct?

        WouterOddBot I'd expect the camera intrinsics to change of the camera zooms in or out. Is that correct?

        We don't currently change the intrinsics; but you are correct, they change since focal length also changes.
        The lens position at which camera is calibrated is written in the eeprom.

        Thanks,
        Jaka

          a month later

          Hi @jakaskerl ,

          We had setAutoFocusLensRange removed now for a while, but the camera still takes quite a bit of time to focus. Not sure if it's working as it should. How much time should the camera need to focus?

          Kind regards,

          Wouter

          Concerning the lens focus and camera intrinsics, how can we compute the correct camera intrinsics given a certain lens position? This really is important to us. If needed I can put that question in a separate topic too.

            WouterOddBot
            You can use OPENCV's calibration procedure to only calibrate the RGB camera.

            • manually set the lens position to the desired one
            • perform calibration using the charuco board and script from opencv
            • you get intrinsics for that lens position
            • repeat for few others

            We are working on fitting a function to the lens position <-> intrinsics relationship. You can wait for the release or do it yourself.

            Thanks,
            Jaka

              Ha thanks great that you're already working on it. We'll wait for the release.

              3 months later

              @jakaskerl Thank you for this helpful thread! Is the lens position <-> intrinsics function available yet for the OAK-D pro poe autofocus camera? It'd be really helpful!

                Joshua
                It's not available because the lenses are not always mounted in the same way due to manufacturing imperfections. We will need a way to interpolate for this information, we are working on it.

                Thanks,
                Jaka

                  Thanks for letting me know Jaka! Do you know if we can expect this interpolation to become available for existing OAK-D pro's by the end of 2024? Just for our own planning.

                    Joshua
                    It's not super high-prio right now so I can't say when it will be implemented.

                    The intrinsics shouldn't change much during normal operation. I guess you can try calibrating for +/- 10 steps of position to see how big the change is.

                    Thanks,
                    Jaka