+
20
-

回答

可使用 Depth Anything V2

import cv2
import numpy as np
import torch
import mediapipe as mp
from PIL import Image
from transformers import pipeline


def load_depth_model():
    device = 0 if torch.cuda.is_available() else -1

    depth_pipe = pipeline(
        task="depth-estimation",
        model="depth-anything/Depth-Anything-V2-Small-hf",
        device=device
    )

    return depth_pipe


def load_person_segmenter():
    mp_selfie = mp.solutions.selfie_segmentation
    return mp_selfie.SelfieSegmentation(model_selection=1)


def get_person_mask(frame_rgb, segmenter):
    result = segmenter.process(frame_rgb)
    mask = result.segmentation_mask > 0.5
    return mask


def main(video_path, output_path="person_depth_da.mp4"):
    cap = cv2.VideoCapture(video_path)

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 30

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    depth_pipe = load_depth_model()
    segmenter = load_person_segmenter()

    while True:
        ret, frame_bgr = cap.read()
        if not ret:
            break

        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

        # Depth Anything 推理
        pil_image = Image.fromarray(frame_rgb)
        result = depth_pipe(pil_image)

        # 得到深度图
        depth = np.array(result["depth"])

        # 保证尺寸一致
        if depth.shape[:2] != frame_rgb.shape[:2]:
            depth = cv2.resize(depth, (frame_rgb.shape[1], frame_rgb.shape[0]))

        # 人物 mask
        mask = get_person_mask(frame_rgb, segmenter)

        # 可视化
        depth_norm = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX)
        depth_norm = depth_norm.astype(np.uint8)
        depth_color = cv2.applyColorMap(depth_norm, cv2.COLORMAP_MAGMA)

        person_depth_color = np.zeros_like(depth_color)
        person_depth_color[mask] = depth_color[mask]

        blended = cv2.addWeighted(frame_bgr, 0.45, person_depth_color, 0.55, 0)

        out.write(blended)
        cv2.imshow("Depth Anything Person Depth", blended)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    cap.release()
    out.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main("input.mp4")

网友回复

我知道答案,我要回答