可使用 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") 网友回复


