import os
|
import cv2
|
import numpy as np
|
from datetime import datetime
|
from ultralytics import YOLO
|
|
|
DEFAULT_YOLOV8S_SEG_MODEL = './models/yolov8s-seg.pt'
|
|
def get_center_distance(box, image_center):
|
x1, y1, x2, y2 = box
|
box_center = ((x1 + x2) / 2, (y1 + y2) / 2)
|
return np.linalg.norm(np.array(box_center) - np.array(image_center))
|
|
def generate_mask_from_first_frame(video_path, filter_type:str = 'dist',mask_type:str="array",mask_file_path:str=None):
|
"""
|
获取视频第一帧的 人像 MASK:
|
filter_type：
|
dist 距离 画面中心最近的人; default
|
box : 画面中面积最大的人;
|
mask_type :
|
array: 直接返回 np.array ; default
|
file：保存成mask_file_path ，返回mask_file_path
|
"""
|
if not os.path.exists(video_path):
|
raise FileNotFoundError(f"视频文件不存在：{video_path}")
|
|
cap = cv2.VideoCapture(video_path)
|
ret, frame = cap.read()
|
cap.release()
|
|
if not ret:
|
raise RuntimeError("无法读取视频第一帧")
|
|
height, width, _ = frame.shape
|
image_center = (width / 2, height / 2)
|
|
|
print(os.path.exists(DEFAULT_YOLOV8S_SEG_MODEL))
|
|
model = YOLO(DEFAULT_YOLOV8S_SEG_MODEL) # 或 yolov8s-seg.pt，体积更大但更准
|
|
# 3. 进行人形检测+分割
|
results = model(frame)[0]
|
|
# 筛选出所有类别为“person”的目标
|
person_masks = []
|
for i, cls in enumerate(results.boxes.cls):
|
if int(cls) == 0: # class 0 是 "person"
|
box = results.boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2]
|
mask = results.masks.data[i].cpu().numpy()
|
|
dist = get_center_distance(box, image_center)
|
box_area = (box[2] - box[0]) * (box[3] - box[1])
|
|
person_masks.append((dist, box_area, mask))
|
|
if not person_masks:
|
return None
|
|
# 4. 取距离中心最近的人形
|
if filter_type=="box":
|
person_masks.sort(key=lambda x: x[1], reverse=True)
|
else:
|
person_masks.sort(key=lambda x: x[0])
|
|
best_mask = person_masks[0][2]
|
best_mask_resized = cv2.resize(best_mask, (width, height), interpolation=cv2.INTER_NEAREST)
|
|
if mask_type == "file":
|
if not mask_file_path:
|
now_str = datetime.now().strftime("%Y%m%d_%H%M%S")
|
mask_file_path =f"./temp/{now_str}.png"
|
mask_img = (best_mask_resized * 255).astype(np.uint8)
|
cv2.imwrite(mask_file_path, mask_img)
|
return mask_file_path
|
else:
|
mask = (best_mask_resized * 255).astype(np.uint8)
|
return mask
|
