From e1a3f40744052164b13dd032dd8876002eadeedf Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Fri, 16 May 2025 15:22:02 -0400 Subject: [PATCH] face style --- .../face_stylization/detector.py | 80 ++++++++++++------- 1 file changed, 51 insertions(+), 29 deletions(-) diff --git a/src/mediapipe_vision/face_stylization/detector.py b/src/mediapipe_vision/face_stylization/detector.py index 8681dba..3661278 100644 --- a/src/mediapipe_vision/face_stylization/detector.py +++ b/src/mediapipe_vision/face_stylization/detector.py @@ -1,45 +1,60 @@ +from typing import Any, List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision +from mediapipe.tasks.python.core.base_options import BaseOptions from ...utils.timing import TimestampProvider +from ..common import BaseDetector -class FaceStylizer: +class FaceStylizer(BaseDetector[torch.Tensor]): """Applies stylization to faces using MediaPipe FaceStylizer.""" def __init__(self, model_path: str): if not model_path: raise ValueError("A valid model_path must be provided.") self.model_path = model_path - self.detector_instance = None + self._detector_instance = None self.current_options = None self._timestamp_provider = None - def _initialize_detector(self): - """Initializes the FaceStylizer.""" - base_options = python.BaseOptions(model_asset_path=self.model_path) + def _create_detector_options(self, base_options: python.BaseOptions, + mode_enum: vision.RunningMode, **kwargs) -> vision.FaceStylizerOptions: + """Create FaceStylizer-specific options. + Note: FaceStylizer only supports IMAGE mode, not VIDEO mode. + """ # FaceStylizerOptions has no specific config other than base options - # Check if FaceStylizer supports VIDEO mode - # According to docs (as of late 2023/early 2024), it seems it only supports IMAGE mode. - # Sticking with IMAGE mode. - options = vision.FaceStylizerOptions( - base_options=base_options - # running_mode=vision.RunningMode.VIDEO # Not supported - ) - self.current_options = () - self._timestamp_provider = TimestampProvider() + # and doesn't support VIDEO mode - we ignore the mode_enum parameter + return vision.FaceStylizerOptions(base_options=base_options) + + def _create_detector_instance(self, options: vision.FaceStylizerOptions) -> vision.FaceStylizer: return vision.FaceStylizer.create_from_options(options) - def stylize(self, image: torch.Tensor) -> torch.Tensor: # Output is IMAGE + def _get_options_tuple(self, running_mode: str = None, delegate: str = None, **kwargs) -> tuple: + # FaceStylizer has no additional options to track + return (delegate,) + + def _process_detection_result(self, stylization_result: Any) -> torch.Tensor: + """Process stylization result to torch tensor.""" + if stylization_result: + # Convert MediaPipe Image back to NumPy (HWC, uint8) + stylized_np = stylization_result.numpy_view() + # Convert to float tensor (HWC, 0-1) + return torch.from_numpy(stylized_np).float().div(255.0) + return None # If no result, we'll handle in the stylize method + + def stylize(self, image: torch.Tensor) -> torch.Tensor: """Applies face stylization to the input image tensor(s).""" if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - if self.detector_instance is None: - self.detector_instance = self._initialize_detector() + # Ensure detector is initialized - always use IMAGE mode + if self._detector_instance is None: + self._initialize_detector("image", "cpu") batch_size = image.shape[0] batch_results = [] @@ -49,20 +64,27 @@ class FaceStylizer: np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) - stylization_result = self.detector_instance.stylize(mp_image) - - output_image_tensor = None - if stylization_result: - # Convert MediaPipe Image back to NumPy (HWC, uint8) - stylized_np = stylization_result.numpy_view() - # Convert to float tensor (HWC, 0-1) - output_image_tensor = torch.from_numpy(stylized_np).float().div(255.0) - else: - # Handle cases where stylization fails (e.g., no face detected) - # Return original image tensor for this batch item + # FaceStylizer only has the stylize method, not detect + stylization_result = self._detector_instance.stylize(mp_image) + output_image_tensor = self._process_detection_result(stylization_result) + + # If processing failed (no face detected), return original image + if output_image_tensor is None: output_image_tensor = img_tensor - + batch_results.append(output_image_tensor) output_batch_tensor = torch.stack(batch_results, dim=0) return output_batch_tensor + + def _initialize_detector(self, running_mode: str, delegate: str): + """Initializes the FaceStylizer.""" + base_options = python.BaseOptions(model_asset_path=self.model_path) + # FaceStylizerOptions has no specific config other than base options + # Check if FaceStylizer supports VIDEO mode + # According to docs (as of late 2023/early 2024), it seems it only supports IMAGE mode. + # Sticking with IMAGE mode. + options = self._create_detector_options(base_options, vision.RunningMode.IMAGE) + self.current_options = () + self._timestamp_provider = TimestampProvider() + self._detector_instance = self._create_detector_instance(options)