Cherokee
eef072e85e
- Replace bbox crop with YOLOv8-seg person segmentation - Full-frame green screen: only person silhouette kept - Top 5 preview images per actor by mask area (IMAGE output) - Accurate time segmentation: only frames where person is visible - No green gaps between segments - Add seg_model_path parameter
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""YOLOv8 person segmentation module."""
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Tuple
|
|
|
|
from core.config import (
|
|
DEFAULT_SEG_MODEL,
|
|
DEFAULT_CONF_THRESHOLD,
|
|
DEFAULT_IOU_THRESHOLD,
|
|
)
|
|
|
|
|
|
class PersonSegmenter:
|
|
"""YOLOv8-seg based person segmenter."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str = DEFAULT_SEG_MODEL,
|
|
device: str = "auto",
|
|
conf_threshold: float = DEFAULT_CONF_THRESHOLD,
|
|
iou_threshold: float = DEFAULT_IOU_THRESHOLD,
|
|
):
|
|
from ultralytics import YOLO
|
|
|
|
self.conf_threshold = conf_threshold
|
|
self.iou_threshold = iou_threshold
|
|
|
|
# Resolve device — same pattern as PersonDetector in detector.py
|
|
if device == "auto":
|
|
try:
|
|
import torch
|
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
|
if torch.cuda.is_available():
|
|
device = "cuda"
|
|
except ImportError:
|
|
device = "cpu"
|
|
|
|
self.device = device
|
|
|
|
# Load model (download if needed — same behavior as yolov8n.pt)
|
|
resolved_path = os.path.expanduser(model_path)
|
|
if not os.path.exists(resolved_path):
|
|
print(f"[PersonSegmenter] Downloading YOLO segmentation model: {model_path}")
|
|
|
|
self.model = YOLO(resolved_path)
|
|
print(f"[PersonSegmenter] Loaded {model_path} on {device}")
|
|
|
|
def detect_masks(self, frame: np.ndarray) -> List[Tuple[int, np.ndarray]]:
|
|
"""
|
|
Detect person segmentation masks in a single frame.
|
|
|
|
Args:
|
|
frame: BGR image (numpy array)
|
|
|
|
Returns:
|
|
list of (track_id_placeholder, mask_2d) tuples for each person detected
|
|
track_id_placeholder is just the detection index (0, 1, 2, ...)
|
|
mask_2d is a boolean numpy array of shape (H, W) where True = person pixel
|
|
"""
|
|
results = self.model(
|
|
frame,
|
|
classes=[0], # person only
|
|
conf=self.conf_threshold,
|
|
iou=self.iou_threshold,
|
|
verbose=False,
|
|
device=self.device,
|
|
)
|
|
|
|
detections = []
|
|
result = results[0]
|
|
|
|
if result.masks is None:
|
|
return detections
|
|
|
|
h, w = frame.shape[:2]
|
|
masks = result.masks.data.cpu().numpy() # (N, mask_h, mask_w)
|
|
|
|
for idx, mask_raw in enumerate(masks):
|
|
# Resize mask to frame size if needed
|
|
mask_resized = np.zeros((h, w), dtype=bool)
|
|
mask_img = (mask_raw * 255).astype(np.uint8)
|
|
mask_img = cv2.resize(mask_img, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
mask_resized = mask_img > 127
|
|
detections.append((idx, mask_resized))
|
|
|
|
return detections
|
|
|
|
def apply_mask(self, frame: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Apply person mask to frame — set non-person pixels to green.
|
|
|
|
Args:
|
|
frame: BGR image (numpy array)
|
|
mask: boolean array (H, W), True = person pixel
|
|
|
|
Returns:
|
|
BGR image where non-person pixels are set to green (0, 255, 0)
|
|
"""
|
|
result = frame.copy()
|
|
result[~mask] = (0, 255, 0)
|
|
return result
|