commit 821a88a3c70e388854628a86118b904189c57bb1 Author: Arkylin Date: Tue Apr 21 15:39:05 2026 +0800 Initial commit: Orbbec software D2C batch converter - d2c_batch.py: interactive profile selection, camera param acquisition, software D2C transformation (Brown-Conrady undistortion + extrinsic projection), batch conversion of uint16 depth PNGs, JET pseudo-color output - utils.py: Orbbec SDK frame conversion helpers (copied from pyorbbecsdk) - README.md: usage guide, parameter format reference, workflow description Co-Authored-By: Claude Sonnet 4.6 diff --git a/README.md b/README.md new file mode 100644 index 0000000..f04f65b --- /dev/null +++ b/README.md @@ -0,0 +1,182 @@ +# Orbbec D2C Batch Converter (软件对齐) + +将奥比中光相机采集的深度图（Depth）批量软件对齐到彩色图（Color）坐标系，输出与 RGB 分辨率一致的对齐深度图。 + +--- + +## 原理 + +软件 D2C（Depth-to-Color）对齐流程： + +``` +深度像素 (u_d, v_d, Z) + ↓ Brown-Conrady 畸变校正 + 反投影 +3D 点（深度相机坐标系）[X, Y, Z] + ↓ 外参旋转 R + 平移 t +3D 点（彩色相机坐标系）[X_c, Y_c, Z_c] + ↓ 彩色内参投影 +彩色像素 (u_c, v_c) → 写入对齐深度图 +``` + +多个深度点映射到同一彩色像素时，保留最近点（小 Z 值覆盖大 Z 值）。 + +--- + +## 环境依赖 + +``` +pyorbbecsdk +numpy +opencv-python +``` + +安装： +```bash +pip install numpy opencv-python +# pyorbbecsdk 参考官方安装说明 +``` + +--- + +## 文件说明 + +| 文件 | 说明 | +|------|------| +| `d2c_batch.py` | 主程序 | +| `camera_params.json` | 相机参数缓存（自动生成，可复用） | +| `utils.py` | 奥比中光 SDK 辅助函数（备用） | + +--- + +## 使用方法 + +### 完整流程（首次使用） + +```bash +python d2c_batch.py +``` + +1. 列出相机所有 **RGB 配置**（分辨率 / 帧率 / 格式），输入序号选择 +2. 列出所有 **Depth 配置**，输入序号选择 +3. 自动启动相机，采集一帧，提取内外参，保存到 `camera_params.json` +4. 输入**深度图目录**路径 +5. 批量转换，结果保存到 `<输入目录>_d2c/` + +### 仅获取相机参数（不转换） + +```bash +python d2c_batch.py --params +``` + +适合先连相机标定参数，稍后离线转换。 + +### 仅批量转换（无需连相机） + +```bash +python d2c_batch.py --convert +``` + +使用已有的 `camera_params.json`，直接输入深度图目录开始转换。适合相机参数已获取、需要反复处理不同数据集的场景。 + +--- + +## 交互示例 + +``` +Found saved camera params: camera_params.json +Use existing params? [y=use existing / n=re-fetch from camera]: n + +=== Phase 1: Camera Parameter Acquisition === + +Available RGB (Color) profiles: + # Resolution FPS Format + ---------------------------------------- + [0 ] 1920x1080 30 MJPG + [1 ] 1280x720 30 MJPG + [2 ] 640x480 30 RGB +Select RGB (Color) profile index: 0 + -> Selected: 1920x1080 @ 30fps MJPG + +Available Depth profiles: + # Resolution FPS Format + ---------------------------------------- + [0 ] 1280x800 30 Y16 + [1 ] 640x400 30 Y16 + [2 ] 320x200 30 Y16 +Select Depth profile index: 1 + -> Selected: 640x400 @ 30fps Y16 + + Starting pipeline... + Pipeline stopped. + Camera params saved to: camera_params.json + +--- Camera Parameters Summary --- + RGB : 1920x1080 @ 30fps MJPG + fx=1382.5 fy=1382.5 cx=959.8 cy=539.4 + Depth : 640x400 @ 30fps Y16 + fx=424.0 fy=424.0 cx=319.5 cy=199.5 + scale=0.001000 mm/unit + Extrinsic t : [-14.82 0.12 0.03] (mm) + +Enter depth images directory: D:\data\depth_raw + +=== Phase 2: Batch D2C Conversion === + Input dir : D:\data\depth_raw + Output dir : D:\data\depth_raw_d2c + Files : 120 + Depth scale: 0.001000 mm/unit + Output size: 1920x1080 (color resolution) + + [ 1/120] 000001.png + [ 2/120] 000002.png + ... + [ 120/120] 000120.png + + Done. 120/120 files converted in 8.3s + Results: D:\data\depth_raw_d2c +``` + +--- + +## 输入 / 输出格式 + +| 项目 | 说明 | +|------|------| +| 输入深度图 | PNG，uint16，单位为传感器原始单位（乘以 `depth_scale` 得 mm） | +| 输出对齐深度图 | PNG，uint16，单位与输入相同，分辨率与所选 RGB 配置一致，目录名 `_d2c` | +| 输出伪彩图 | PNG，uint8 BGR，JET colormap，无效像素为纯黑，目录名 `_d2c_vis` | + +--- + +## camera_params.json 格式 + +```json +{ + "color": { + "width": 1920, "height": 1080, "fps": 30, "format": "MJPG", + "intrinsic": { "fx": 1382.5, "fy": 1382.5, "cx": 959.8, "cy": 539.4, "width": 1920, "height": 1080 }, + "distortion": { "k1": -0.055, "k2": 0.071, "k3": 0.0, "k4": 0.0, "k5": 0.0, "k6": 0.0, "p1": 0.0, "p2": 0.0 } + }, + "depth": { + "width": 640, "height": 400, "fps": 30, "format": "Y16", + "scale": 0.001, + "intrinsic": { "fx": 424.0, "fy": 424.0, "cx": 319.5, "cy": 199.5, "width": 640, "height": 400 }, + "distortion": { "k1": 0.0, ... } + }, + "extrinsic": { + "rot": [1,0,0, 0,1,0, 0,0,1], + "transform": [-14.82, 0.12, 0.03] + } +} +``` + +`extrinsic.rot` 为行优先展平的 3×3 旋转矩阵，`extrinsic.transform` 为平移向量（单位 mm），方向为深度相机坐标系 → 彩色相机坐标系。 + +--- + +## 注意事项 + +- 获取相机参数时需要相机**实际连接**；批量转换时不需要。 +- 更换相机或修改分辨率后需重新获取参数（运行时选 `n` 或使用 `--params`）。 +- 输入深度图分辨率须与获取参数时选择的 Depth 分辨率一致。 +- 对齐深度图的有效区域受深度传感器视场角限制，彩色图边缘区域可能无深度值（值为 0）。 diff --git a/d2c_batch.py b/d2c_batch.py new file mode 100644 index 0000000..b14abe4 --- /dev/null +++ b/d2c_batch.py @@ -0,0 +1,455 @@ +#!/usr/bin/env python3 +""" +Orbbec D2C Batch Converter (Software) + +Usage: + python d2c_batch.py # Interactive: select profiles, fetch params, convert + python d2c_batch.py --params # Only fetch/update camera params (no conversion) + python d2c_batch.py --convert # Only convert using saved camera_params.json + +Workflow: + 1. Select RGB profile (resolution / fps / format) + 2. Select Depth profile (resolution / fps / format) + 3. Camera intrinsics + extrinsics are fetched and saved to camera_params.json + 4. Enter depth image directory + 5. Batch-convert all PNG depth images via software D2C + -> Output saved to _d2c/ +""" + +import os +import sys +import json +import time +import numpy as np +import cv2 +from pathlib import Path + +PARAMS_FILE = "camera_params.json" + + +# --------------------------------------------------------------------------- +# Camera parameter acquisition +# --------------------------------------------------------------------------- + +def _list_video_profiles(profile_list): + """Return list of (index, VideoStreamProfile) for all video profiles.""" + results = [] + count = profile_list.get_count() + for i in range(count): + p = profile_list.get_stream_profile_by_index(i) + vp = p.as_video_stream_profile() + if vp is None: + continue + results.append((i, vp)) + return results + + +def _print_profiles(profiles, sensor_name): + print(f"\nAvailable {sensor_name} profiles:") + print(f" {'#':<4} {'Resolution':<14} {'FPS':<6} {'Format'}") + print(f" {'-'*40}") + for idx, vp in profiles: + fmt = str(vp.get_format()).split(".")[-1] + print(f" [{idx:<2}] {vp.get_width()}x{vp.get_height():<8} {vp.get_fps():<6} {fmt}") + + +def _select_profile(profile_list, sensor_name): + """Interactive profile selection. Returns the selected VideoStreamProfile.""" + profiles = _list_video_profiles(profile_list) + if not profiles: + print(f" No {sensor_name} profiles found!") + return None + + _print_profiles(profiles, sensor_name) + indices = [i for i, _ in profiles] + + while True: + raw = input(f"Select {sensor_name} profile index: ").strip() + try: + choice = int(raw) + if choice in indices: + selected = next(vp for i, vp in profiles if i == choice) + fmt = str(selected.get_format()).split(".")[-1] + print(f" -> Selected: {selected.get_width()}x{selected.get_height()} @ {selected.get_fps()}fps {fmt}") + return selected + except ValueError: + pass + print(f" Invalid. Choose from: {indices}") + + +def _intrinsic_to_dict(intr): + return { + "fx": float(intr.fx), "fy": float(intr.fy), + "cx": float(intr.cx), "cy": float(intr.cy), + "width": int(intr.width), "height": int(intr.height), + } + + +def _distortion_to_dict(dist): + return { + "k1": float(dist.k1), "k2": float(dist.k2), + "k3": float(dist.k3), "k4": float(dist.k4), + "k5": float(dist.k5), "k6": float(dist.k6), + "p1": float(dist.p1), "p2": float(dist.p2), + } + + +def fetch_camera_params(): + """ + Interactively select color + depth profiles, start the pipeline briefly to + capture one frameset, extract intrinsics/extrinsics, and return as a dict. + """ + from pyorbbecsdk import Pipeline, Config, OBSensorType + + print("\n=== Phase 1: Camera Parameter Acquisition ===") + pipeline = Pipeline() + config = Config() + + # --- Color profile --- + try: + color_list = pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR) + except Exception as e: + print(f" Cannot get color profiles: {e}") + return None + + color_profile = _select_profile(color_list, "RGB (Color)") + if color_profile is None: + return None + + # --- Depth profile --- + try: + depth_list = pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR) + except Exception as e: + print(f" Cannot get depth profiles: {e}") + return None + + depth_profile = _select_profile(depth_list, "Depth") + if depth_profile is None: + return None + + config.enable_stream(color_profile) + config.enable_stream(depth_profile) + + print("\n Starting pipeline...") + try: + pipeline.start(config) + except Exception as e: + print(f" Failed to start pipeline: {e}") + return None + + # --- Wait for first valid frameset --- + frames = None + deadline = time.time() + 10.0 + while time.time() < deadline: + f = pipeline.wait_for_frames(200) + if f and f.get_color_frame() and f.get_depth_frame(): + frames = f + break + + if frames is None: + print(" Timeout: no frameset received within 10 s.") + pipeline.stop() + return None + + color_frame = frames.get_color_frame() + depth_frame = frames.get_depth_frame() + + # --- Extract profile-specific intrinsics from actual frames --- + color_vp = color_frame.get_stream_profile().as_video_stream_profile() + depth_vp = depth_frame.get_stream_profile().as_video_stream_profile() + + color_intr = color_vp.get_intrinsic() + color_dist = color_vp.get_distortion() + depth_intr = depth_vp.get_intrinsic() + depth_dist = depth_vp.get_distortion() + + # Extrinsic: depth camera -> color camera + extrinsic = depth_vp.get_extrinsic_to(color_vp) + + depth_scale = depth_frame.get_depth_scale() # mm per raw unit + + pipeline.stop() + print(" Pipeline stopped.") + + params = { + "color": { + "width": color_frame.get_width(), + "height": color_frame.get_height(), + "fps": color_profile.get_fps(), + "format": str(color_profile.get_format()).split(".")[-1], + "intrinsic": _intrinsic_to_dict(color_intr), + "distortion": _distortion_to_dict(color_dist), + }, + "depth": { + "width": depth_frame.get_width(), + "height": depth_frame.get_height(), + "fps": depth_profile.get_fps(), + "format": str(depth_profile.get_format()).split(".")[-1], + "scale": float(depth_scale), + "intrinsic": _intrinsic_to_dict(depth_intr), + "distortion": _distortion_to_dict(depth_dist), + }, + # rot: 9-element flat array (row-major 3x3), transform: 3-element translation (mm) + "extrinsic": { + "rot": extrinsic.rot.tolist(), + "transform": extrinsic.transform.tolist(), + }, + } + return params + + +# --------------------------------------------------------------------------- +# Software D2C transformation +# --------------------------------------------------------------------------- + +def _undistort_points(u, v, intr, dist): + """Brown-Conrady lens undistortion (returns undistorted normalised coords).""" + fx, fy = intr["fx"], intr["fy"] + cx, cy = intr["cx"], intr["cy"] + k1, k2, k3 = dist["k1"], dist["k2"], dist["k3"] + p1, p2 = dist["p1"], dist["p2"] + + x = (u - cx) / fx + y = (v - cy) / fy + + r2 = x * x + y * y + radial = 1.0 + k1 * r2 + k2 * r2**2 + k3 * r2**3 + x_u = x * radial + 2.0 * p1 * x * y + p2 * (r2 + 2.0 * x * x) + y_u = y * radial + p1 * (r2 + 2.0 * y * y) + 2.0 * p2 * x * y + return x_u, y_u + + +def software_d2c(depth_img, params): + """ + Transform a uint16 depth image (depth camera space) to an aligned uint16 + depth image in color camera space. + + Parameters + ---------- + depth_img : np.ndarray (H_d x W_d, uint16) raw depth in sensor units + params : dict camera_params.json content + + Returns + ------- + aligned : np.ndarray (H_c x W_c, uint16) aligned depth in sensor units + """ + depth_intr = params["depth"]["intrinsic"] + depth_dist = params["depth"]["distortion"] + color_intr = params["color"]["intrinsic"] + ext = params["extrinsic"] + + color_w = color_intr["width"] + color_h = color_intr["height"] + fx_c, fy_c = color_intr["fx"], color_intr["fy"] + cx_c, cy_c = color_intr["cx"], color_intr["cy"] + + dh, dw = depth_img.shape + depth_scale = params["depth"]["scale"] + + # Pixel grid + u_d = np.arange(dw, dtype=np.float32) + v_d = np.arange(dh, dtype=np.float32) + u_d, v_d = np.meshgrid(u_d, v_d) # (dh, dw) + + # Depth in mm (float) + Z = depth_img.astype(np.float32) * depth_scale + valid = Z > 0 + + # Undistort + unproject to 3-D (depth camera space, mm) + x_u, y_u = _undistort_points(u_d, v_d, depth_intr, depth_dist) + X = x_u * Z # (dh, dw) + Y = y_u * Z + + # Extrinsic: R (3x3) and t (3,) in mm + R = np.asarray(ext["rot"], dtype=np.float64).reshape(3, 3) + t = np.asarray(ext["transform"], dtype=np.float64) + + # Vectorised transform [X_c, Y_c, Z_c] = R @ [X, Y, Z]^T + t + pts = np.stack([X.ravel(), Y.ravel(), Z.ravel()], axis=0).astype(np.float64) # (3, N) + pts_c = R @ pts + t[:, np.newaxis] # (3, N) + + Xc = pts_c[0].reshape(dh, dw).astype(np.float32) + Yc = pts_c[1].reshape(dh, dw).astype(np.float32) + Zc = pts_c[2].reshape(dh, dw).astype(np.float32) + + # Project onto color image plane + valid_c = valid & (Zc > 0) + u_c = np.where(valid_c, fx_c * Xc / Zc + cx_c, -1.0).astype(np.float32) + v_c = np.where(valid_c, fy_c * Yc / Zc + cy_c, -1.0).astype(np.float32) + + u_ci = np.round(u_c).astype(np.int32) + v_ci = np.round(v_c).astype(np.int32) + + in_bounds = (u_ci >= 0) & (u_ci < color_w) & (v_ci >= 0) & (v_ci < color_h) + mask = valid_c & in_bounds + + u_vals = u_ci[mask] + v_vals = v_ci[mask] + z_vals = depth_img[mask] # keep original uint16 sensor units + + # Write far pixels first so nearer pixels overwrite (z small = close) + order = np.argsort(z_vals)[::-1] + u_vals = u_vals[order] + v_vals = v_vals[order] + z_vals = z_vals[order] + + aligned = np.zeros((color_h, color_w), dtype=np.uint16) + aligned[v_vals, u_vals] = z_vals + return aligned + + +# --------------------------------------------------------------------------- +# Batch conversion +# --------------------------------------------------------------------------- + +def depth_to_colormap(aligned, min_depth_mm=200, max_depth_mm=5000, depth_scale=1.0): + """Convert uint16 aligned depth to a JET pseudo-color BGR image.""" + depth_mm = aligned.astype(np.float32) * depth_scale + valid = (depth_mm > min_depth_mm) & (depth_mm < max_depth_mm) + norm = np.zeros_like(depth_mm) + norm[valid] = (depth_mm[valid] - min_depth_mm) / (max_depth_mm - min_depth_mm) + norm = np.clip(norm, 0.0, 1.0) + gray = (norm * 255).astype(np.uint8) + colormap = cv2.applyColorMap(gray, cv2.COLORMAP_JET) + colormap[~valid] = 0 # black for invalid pixels + return colormap + + +def batch_convert(params, depth_dir_str): + depth_dir = Path(depth_dir_str.strip().strip('"\'')) + if not depth_dir.exists(): + print(f" Directory not found: {depth_dir}") + return + + png_files = sorted(depth_dir.glob("*.png")) + sorted(depth_dir.glob("*.PNG")) + png_files = sorted(set(png_files)) + if not png_files: + print(f" No PNG files found in: {depth_dir}") + return + + out_dir = depth_dir.parent / (depth_dir.name + "_d2c") + vis_dir = depth_dir.parent / (depth_dir.name + "_d2c_vis") + out_dir.mkdir(exist_ok=True) + vis_dir.mkdir(exist_ok=True) + + color_w = params["color"]["intrinsic"]["width"] + color_h = params["color"]["intrinsic"]["height"] + depth_scale = params["depth"]["scale"] + + print(f"\n=== Phase 2: Batch D2C Conversion ===") + print(f" Input dir : {depth_dir}") + print(f" Aligned depth: {out_dir}") + print(f" Pseudo-color : {vis_dir}") + print(f" Files : {len(png_files)}") + print(f" Depth scale : {depth_scale:.6f} mm/unit") + print(f" Output size : {color_w}x{color_h} (color resolution)") + print() + + t0 = time.time() + ok = 0 + for i, fpath in enumerate(png_files, 1): + depth_img = cv2.imread(str(fpath), cv2.IMREAD_UNCHANGED) + if depth_img is None: + print(f" [{i:>4}/{len(png_files)}] SKIP (unreadable): {fpath.name}") + continue + if depth_img.ndim != 2: + depth_img = depth_img[:, :, 0] + depth_img = depth_img.astype(np.uint16) + + aligned = software_d2c(depth_img, params) + + # Save aligned uint16 depth + cv2.imwrite(str(out_dir / fpath.name), aligned) + + # Save pseudo-color visualization + colormap = depth_to_colormap(aligned, depth_scale=depth_scale) + vis_name = fpath.stem + "_vis.png" + cv2.imwrite(str(vis_dir / vis_name), colormap) + + ok += 1 + print(f" [{i:>4}/{len(png_files)}] {fpath.name}") + + elapsed = time.time() - t0 + print(f"\n Done. {ok}/{len(png_files)} files converted in {elapsed:.1f}s") + print(f" Aligned depth : {out_dir}") + print(f" Pseudo-color : {vis_dir}") + + +# --------------------------------------------------------------------------- +# Parameter display +# --------------------------------------------------------------------------- + +def print_params_summary(params): + c = params["color"] + d = params["depth"] + ci = c["intrinsic"] + di = d["intrinsic"] + ext = params["extrinsic"] + R = np.array(ext["rot"]).reshape(3, 3) + t = np.array(ext["transform"]) + + print("\n--- Camera Parameters Summary ---") + print(f" RGB : {c['width']}x{c['height']} @ {c['fps']}fps {c['format']}") + print(f" fx={ci['fx']:.4f} fy={ci['fy']:.4f} cx={ci['cx']:.4f} cy={ci['cy']:.4f}") + print(f" Depth : {d['width']}x{d['height']} @ {d['fps']}fps {d['format']}") + print(f" fx={di['fx']:.4f} fy={di['fy']:.4f} cx={di['cx']:.4f} cy={di['cy']:.4f}") + print(f" scale={d['scale']:.6f} mm/unit") + print(f" Extrinsic R :\n {R[0]}\n {R[1]}\n {R[2]}") + print(f" Extrinsic t : {t} (mm)") + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def main(): + only_params = "--params" in sys.argv + only_convert = "--convert" in sys.argv + + params = None + + # ---- Load or fetch camera parameters ---- + if not only_convert: + if os.path.exists(PARAMS_FILE) and not only_params: + print(f"Found saved camera params: {PARAMS_FILE}") + ans = input("Use existing params? [y=use existing / n=re-fetch from camera]: ").strip().lower() + if ans == "y": + with open(PARAMS_FILE) as f: + params = json.load(f) + print(" Loaded existing params.") + + if params is None: + params = fetch_camera_params() + if params is None: + print("Failed to fetch camera parameters. Exiting.") + sys.exit(1) + with open(PARAMS_FILE, "w") as f: + json.dump(params, f, indent=2) + print(f"\n Camera params saved to: {PARAMS_FILE}") + + else: + # --convert: must have existing params file + if not os.path.exists(PARAMS_FILE): + print(f"No {PARAMS_FILE} found. Run without --convert first to fetch camera params.") + sys.exit(1) + with open(PARAMS_FILE) as f: + params = json.load(f) + print(f"Loaded camera params from: {PARAMS_FILE}") + + print_params_summary(params) + + if only_params: + print("\n--params mode: done.") + return + + # ---- Batch conversion ---- + depth_dir = input("\nEnter depth images directory: ").strip() + if not depth_dir: + print("No directory entered. Exiting.") + return + + batch_convert(params, depth_dir) + + +if __name__ == "__main__": + main() diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..5977760 --- /dev/null +++ b/utils.py @@ -0,0 +1,140 @@ +# ****************************************************************************** +# Copyright (c) 2024 Orbbec 3D Technology, Inc +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ****************************************************************************** +from typing import Union, Any, Optional + +import cv2 +import numpy as np + +from pyorbbecsdk import FormatConvertFilter, VideoFrame, Device +from pyorbbecsdk import OBFormat, OBConvertFormat, OBSensorType + +def is_astra_mini_device(vid: int, pid: int) -> bool: + if (vid == 0x2bc5) and (pid == 0x069d or pid == 0x069d or pid ==0x065b or pid == 0x065e): + return True + return False + +def is_lidar_device(device: Device) -> bool: + sensor_list = device.get_sensor_list() + count = sensor_list.get_count() + for index in range(count): + sensor_type = sensor_list.get_sensor_by_index(index).get_type() + if sensor_type == OBSensorType.LIDAR_SENSOR: + return True + return False + +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