保存图片快照

api.py

@@ -2,6 +2,7 @@ import os
 import time
 import threading
 
+import cv2
 import numpy as np
 from fastapi import FastAPI, HTTPException
 import uvicorn
@@ -11,14 +12,17 @@ from depth_common import (
     TemporalFilter,
     compute_roi_bounds,
     extract_depth_data,
+    find_nearest_point,
     init_depth_pipeline,
     nearest_distance_in_roi,
 )
+from utils import frame_to_bgr_image
 
 # 采样参数
 SAMPLE_COUNT = 10
 FRAME_TIMEOUT_MS = 200
 SAMPLE_TIMEOUT_SEC = 8
+MAX_SAVED_IMAGES = int(os.getenv("MAX_SAVED_IMAGES", "1000"))
 # 从环境变量加载测量配置
 SETTINGS = Settings.from_env()
 
@@ -60,6 +64,7 @@ def _measure_once():
     frames = _pipeline.wait_for_frames(FRAME_TIMEOUT_MS)
     if frames is None:
         return None
+    color_frame = frames.get_color_frame()
     depth_frame = frames.get_depth_frame()
     depth_data = extract_depth_data(depth_frame, SETTINGS, _temporal_filter)
     if depth_data is None:
@@ -67,9 +72,102 @@ def _measure_once():
     bounds = compute_roi_bounds(depth_data, _depth_intrinsics, SETTINGS)
     if bounds is None:
         return None
-    x_start, x_end, y_start, y_end, _ = bounds
+    x_start, x_end, y_start, y_end, center_distance = bounds
     roi = depth_data[y_start:y_end, x_start:x_end]
-    return nearest_distance_in_roi(roi, SETTINGS)
+    nearest_distance = nearest_distance_in_roi(roi, SETTINGS)
+    if nearest_distance is None:
+        return None
+    return {
+        "nearest_distance": nearest_distance,
+        "color_frame": color_frame,
+        "depth_data": depth_data,
+        "bounds": bounds,
+        "center_distance": center_distance,
+    }
+
+
+def _save_current_sample_images(sample):
+    save_image_dir = os.path.join(os.getcwd(), "sample_images")
+    os.makedirs(save_image_dir, exist_ok=True)
+    now = time.localtime()
+    time_str = time.strftime("%Y%m%d_%H%M%S", now)
+    millis = int((time.time() % 1) * 1000)
+    timestamp = f"{time_str}_{millis:03d}"
+
+    color_frame = sample.get("color_frame")
+    if color_frame is not None:
+        color_image = frame_to_bgr_image(color_frame)
+        if color_image is not None:
+            color_height, color_width = color_image.shape[:2]
+            color_file = os.path.join(
+                save_image_dir,
+                f"color_{color_width}x{color_height}_{timestamp}.png",
+            )
+            cv2.imwrite(color_file, color_image)
+
+    depth_data = sample["depth_data"]
+    x_start, x_end, y_start, y_end, center_distance = sample["bounds"]
+    nearest_distance = sample["nearest_distance"]
+    roi = depth_data[y_start:y_end, x_start:x_end]
+    nearest_point = find_nearest_point(roi, x_start, y_start, SETTINGS, nearest_distance)
+
+    depth_image = cv2.normalize(depth_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+    depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)
+    cv2.rectangle(
+        depth_image,
+        (x_start, y_start),
+        (x_end - 1, y_end - 1),
+        (0, 255, 0),
+        2,
+    )
+    if nearest_point is not None:
+        cv2.circle(depth_image, nearest_point, 4, (0, 0, 0), -1)
+        cv2.circle(depth_image, nearest_point, 6, (0, 255, 255), 2)
+
+    cv2.putText(
+        depth_image,
+        f"nearest: {nearest_distance} mm",
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        0.8,
+        (255, 255, 255),
+        2,
+        cv2.LINE_AA,
+    )
+    cv2.putText(
+        depth_image,
+        f"center: {int(center_distance)} mm",
+        (10, 60),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        0.8,
+        (255, 255, 255),
+        2,
+        cv2.LINE_AA,
+    )
+
+    depth_h, depth_w = depth_image.shape[:2]
+    depth_file = os.path.join(
+        save_image_dir,
+        f"depth_annotated_{depth_w}x{depth_h}_{timestamp}.png",
+    )
+    cv2.imwrite(depth_file, depth_image)
+    _prune_saved_images(save_image_dir, MAX_SAVED_IMAGES)
+
+
+def _prune_saved_images(save_dir, max_images):
+    png_files = [
+        os.path.join(save_dir, name)
+        for name in os.listdir(save_dir)
+        if name.lower().endswith(".png")
+    ]
+    if len(png_files) <= max_images:
+        return
+    png_files.sort(key=os.path.getmtime)
+    for file_path in png_files[: len(png_files) - max_images]:
+        try:
+            os.remove(file_path)
+        except OSError:
+            pass
 
 
 @app.on_event("startup")
@@ -89,11 +187,33 @@ def get_height():
     # 采集多次样本并返回中位数高度
     start_time = time.time()
     samples = []
+    first_valid_sample = None
+    first_color_frame = None
     with _lock:
         while len(samples) < SAMPLE_COUNT and (time.time() - start_time) < SAMPLE_TIMEOUT_SEC:
-            value = _measure_once()
-            if value is not None:
-                samples.append(value)
+            sample = _measure_once()
+            if sample is not None:
+                samples.append(sample["nearest_distance"])
+                if first_valid_sample is None:
+                    first_valid_sample = sample
+                if first_color_frame is None and sample.get("color_frame") is not None:
+                    first_color_frame = sample.get("color_frame")
+
+        # If no color frame arrived during valid depth sampling, try a few extra pulls.
+        if first_color_frame is None:
+            for _ in range(5):
+                frames = _pipeline.wait_for_frames(FRAME_TIMEOUT_MS)
+                if frames is None:
+                    continue
+                color_frame = frames.get_color_frame()
+                if color_frame is not None:
+                    first_color_frame = color_frame
+                    break
+
+        if first_valid_sample is not None:
+            if first_color_frame is not None:
+                first_valid_sample["color_frame"] = first_color_frame
+            _save_current_sample_images(first_valid_sample)
     if len(samples) < SAMPLE_COUNT:
         raise HTTPException(status_code=503, detail="Insufficient valid samples from depth camera")
     median_value = int(np.median(np.array(samples, dtype=np.int32)))

depth_common.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 import cv2
 import numpy as np
 
-from pyorbbecsdk import Config, Pipeline, OBSensorType, OBFormat
+from pyorbbecsdk import Config, Pipeline, OBSensorType, OBFormat, OBError
 
 
 def _get_env_int(name, default):
@@ -77,6 +77,18 @@ def init_depth_pipeline():
         raise RuntimeError("default depth profile is empty")
     depth_intrinsics = depth_profile.get_intrinsic()
     config.enable_stream(depth_profile)
+
+    # Optionally enable color stream so callers can access color frames when available.
+    try:
+        color_profile_list = pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
+        if color_profile_list is not None:
+            color_profile = color_profile_list.get_default_video_stream_profile()
+            if color_profile is not None:
+                config.enable_stream(color_profile)
+    except OBError:
+        # Some devices do not provide color sensors; keep depth-only behavior.
+        pass
+
     pipeline.start(config)
     return pipeline, depth_intrinsics, depth_profile
 

examples/quick_start.py

@@ -0,0 +1,76 @@
+import cv2
+import numpy as np
+import time
+
+from pyorbbecsdk import *
+from utils import frame_to_bgr_image
+
+ESC_KEY = 27
+MIN_DEPTH = 20  # 20mm
+MAX_DEPTH = 10000  # 10000mm
+
+
+def main():
+    pipeline = Pipeline()
+
+    pipeline.start()
+    print("Pipeline started successfully. Press 'q' or ESC to exit.")
+
+    # Set window size
+    window_width = 1280
+    window_height = 720
+    cv2.namedWindow("QuickStart Viewer", cv2.WINDOW_NORMAL)
+    cv2.resizeWindow("QuickStart Viewer", window_width, window_height)
+
+    while True:
+        try:
+            frames = pipeline.wait_for_frames(100)
+            if frames is None:
+                continue
+
+            # Get color frame
+            color_frame = frames.get_color_frame()
+            if color_frame is None:
+                continue
+            color_image = frame_to_bgr_image(color_frame)
+
+            # Get depth frame
+            depth_frame = frames.get_depth_frame()
+            if depth_frame is None:
+                continue
+            if depth_frame.get_format() != OBFormat.Y16:
+                print("Depth format is not Y16")
+                continue
+
+            # Process depth data
+            width = depth_frame.get_width()
+            height = depth_frame.get_height()
+            scale = depth_frame.get_depth_scale()
+
+            depth_data = np.frombuffer(depth_frame.get_data(), dtype=np.uint16).reshape((height, width))
+            depth_data = depth_data.astype(np.float32) * scale
+            depth_data = np.where((depth_data > MIN_DEPTH) & (depth_data < MAX_DEPTH), depth_data, 0).astype(np.uint16)
+
+            # Create depth visualization
+            depth_image = cv2.normalize(depth_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+            depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)
+
+            # Resize and combine images
+            color_image_resized = cv2.resize(color_image, (window_width // 2, window_height))
+            depth_image_resized = cv2.resize(depth_image, (window_width // 2, window_height))
+            combined_image = np.hstack((color_image_resized, depth_image_resized))
+
+            cv2.imshow("QuickStart Viewer", combined_image)
+
+            if cv2.waitKey(1) in [ord('q'), ESC_KEY]:
+                break
+        except KeyboardInterrupt:
+            break
+
+    cv2.destroyAllWindows()
+    pipeline.stop()
+    print("Pipeline stopped and all windows closed.")
+
+
+if __name__ == "__main__":
+    main()

examples/save_image_to_disk.py

@@ -0,0 +1,91 @@
+import os
+
+import cv2
+import numpy as np
+
+from pyorbbecsdk import *
+from utils import frame_to_bgr_image
+
+
+def save_depth_frame(frame: DepthFrame, index):
+    if frame is None:
+        return
+    width = frame.get_width()
+    height = frame.get_height()
+    timestamp = frame.get_timestamp()
+    scale = frame.get_depth_scale()
+    depth_format = frame.get_format()
+    if depth_format != OBFormat.Y16:
+        print("depth format is not Y16")
+        return
+    data = np.frombuffer(frame.get_data(), dtype=np.uint16)
+    data = data.reshape((height, width))
+    data = data.astype(np.float32) * scale
+    data = data.astype(np.uint16)
+    save_image_dir = os.path.join(os.getcwd(), "depth_images")
+    if not os.path.exists(save_image_dir):
+        os.mkdir(save_image_dir)
+    raw_filename = save_image_dir + "/depth_{}x{}_{}_{}.raw".format(width, height, index, timestamp)
+    data.tofile(raw_filename)
+
+
+def save_color_frame(frame: ColorFrame, index):
+    if frame is None:
+        return
+    width = frame.get_width()
+    height = frame.get_height()
+    timestamp = frame.get_timestamp()
+    save_image_dir = os.path.join(os.getcwd(), "color_images")
+    if not os.path.exists(save_image_dir):
+        os.mkdir(save_image_dir)
+    filename = save_image_dir + "/color_{}x{}_{}_{}.png".format(width, height, index, timestamp)
+    image = frame_to_bgr_image(frame)
+    if image is None:
+        print("failed to convert frame to image")
+        return
+    cv2.imwrite(filename, image)
+
+
+def main():
+    pipeline = Pipeline()
+    config = Config()
+    saved_color_cnt: int = 0
+    saved_depth_cnt: int = 0
+    has_color_sensor = False
+    try:
+        profile_list = pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
+        if profile_list is not None:
+            color_profile: VideoStreamProfile = profile_list.get_default_video_stream_profile()
+            config.enable_stream(color_profile)
+            has_color_sensor = True
+    except OBError as e:
+        print(e)
+    depth_profile_list = pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR)
+    if depth_profile_list is not None:
+        depth_profile = depth_profile_list.get_default_video_stream_profile()
+        config.enable_stream(depth_profile)
+    pipeline.start(config)
+    while True:
+        try:
+            frames = pipeline.wait_for_frames(100)
+            if frames is None:
+                continue
+            if has_color_sensor:
+                if saved_color_cnt >= 5 and saved_depth_cnt >= 5:
+                    break
+            elif saved_depth_cnt >= 5:
+                break
+            color_frame = frames.get_color_frame()
+            if color_frame is not None and saved_color_cnt < 5:
+                save_color_frame(color_frame, saved_color_cnt)
+                saved_color_cnt += 1
+            depth_frame = frames.get_depth_frame()
+            if depth_frame is not None and saved_depth_cnt < 5:
+                save_depth_frame(depth_frame, saved_depth_cnt)
+                saved_depth_cnt += 1
+        except KeyboardInterrupt:
+            break
+
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -0,0 +1,112 @@
+from typing import Union, Any, Optional
+
+import cv2
+import numpy as np
+
+from pyorbbecsdk import FormatConvertFilter, VideoFrame
+from pyorbbecsdk import OBFormat, OBConvertFormat
+
+
+def yuyv_to_bgr(frame: np.ndarray, width: int, height: int) -> np.ndarray:
+    yuyv = frame.reshape((height, width, 2))
+    bgr_image = cv2.cvtColor(yuyv, cv2.COLOR_YUV2BGR_YUY2)
+    return bgr_image
+
+
+def uyvy_to_bgr(frame: np.ndarray, width: int, height: int) -> np.ndarray:
+    uyvy = frame.reshape((height, width, 2))
+    bgr_image = cv2.cvtColor(uyvy, cv2.COLOR_YUV2BGR_UYVY)
+    return bgr_image
+
+
+def i420_to_bgr(frame: np.ndarray, width: int, height: int) -> np.ndarray:
+    y = frame[0:height, :]
+    u = frame[height:height + height // 4].reshape(height // 2, width // 2)
+    v = frame[height + height // 4:].reshape(height // 2, width // 2)
+    yuv_image = cv2.merge([y, u, v])
+    bgr_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_I420)
+    return bgr_image
+
+
+def nv21_to_bgr(frame: np.ndarray, width: int, height: int) -> np.ndarray:
+    y = frame[0:height, :]
+    uv = frame[height:height + height // 2].reshape(height // 2, width)
+    yuv_image = cv2.merge([y, uv])
+    bgr_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_NV21)
+    return bgr_image
+
+
+def nv12_to_bgr(frame: np.ndarray, width: int, height: int) -> np.ndarray:
+    y = frame[0:height, :]
+    uv = frame[height:height + height // 2].reshape(height // 2, width)
+    yuv_image = cv2.merge([y, uv])
+    bgr_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_NV12)
+    return bgr_image
+
+
+def determine_convert_format(frame: VideoFrame):
+    if frame.get_format() == OBFormat.I420:
+        return OBConvertFormat.I420_TO_RGB888
+    elif frame.get_format() == OBFormat.MJPG:
+        return OBConvertFormat.MJPG_TO_RGB888
+    elif frame.get_format() == OBFormat.YUYV:
+        return OBConvertFormat.YUYV_TO_RGB888
+    elif frame.get_format() == OBFormat.NV21:
+        return OBConvertFormat.NV21_TO_RGB888
+    elif frame.get_format() == OBFormat.NV12:
+        return OBConvertFormat.NV12_TO_RGB888
+    elif frame.get_format() == OBFormat.UYVY:
+        return OBConvertFormat.UYVY_TO_RGB888
+    else:
+        return None
+
+
+def frame_to_rgb_frame(frame: VideoFrame) -> Union[Optional[VideoFrame], Any]:
+    if frame.get_format() == OBFormat.RGB:
+        return frame
+    convert_format = determine_convert_format(frame)
+    if convert_format is None:
+        print("Unsupported format")
+        return None
+    print("covert format: {}".format(convert_format))
+    convert_filter = FormatConvertFilter()
+    convert_filter.set_format_convert_format(convert_format)
+    rgb_frame = convert_filter.process(frame)
+    if rgb_frame is None:
+        print("Convert {} to RGB failed".format(frame.get_format()))
+    return rgb_frame
+
+
+def frame_to_bgr_image(frame: VideoFrame) -> Union[Optional[np.array], Any]:
+    width = frame.get_width()
+    height = frame.get_height()
+    color_format = frame.get_format()
+    data = np.asanyarray(frame.get_data())
+    image = np.zeros((height, width, 3), dtype=np.uint8)
+    if color_format == OBFormat.RGB:
+        image = np.resize(data, (height, width, 3))
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    elif color_format == OBFormat.BGR:
+        image = np.resize(data, (height, width, 3))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    elif color_format == OBFormat.YUYV:
+        image = np.resize(data, (height, width, 2))
+        image = cv2.cvtColor(image, cv2.COLOR_YUV2BGR_YUYV)
+    elif color_format == OBFormat.MJPG:
+        image = cv2.imdecode(data, cv2.IMREAD_COLOR)
+    elif color_format == OBFormat.I420:
+        image = i420_to_bgr(data, width, height)
+        return image
+    elif color_format == OBFormat.NV12:
+        image = nv12_to_bgr(data, width, height)
+        return image
+    elif color_format == OBFormat.NV21:
+        image = nv21_to_bgr(data, width, height)
+        return image
+    elif color_format == OBFormat.UYVY:
+        image = np.resize(data, (height, width, 2))
+        image = cv2.cvtColor(image, cv2.COLOR_YUV2BGR_UYVY)
+    else:
+        print("Unsupported color format: {}".format(color_format))
+        return None
+    return image