Spaces:
Sleeping
Sleeping
| from PIL import Image | |
| from kaggle_gpu_server.engine.plugin_base import ModelPlugin, PluginCapability, ModelCategory | |
| class DWPosePlugin(ModelPlugin): | |
| name = "dwpose" | |
| model_id = "yzd-v/DWPose" | |
| capability = PluginCapability.POSE_ESTIMATION | |
| category = ModelCategory.LIGHTWEIGHT | |
| vram_estimate_mb = 300 | |
| version = "1.0.0" | |
| description = "Estimates pose using DWPose and returns skeleton visualization" | |
| def __init__(self): | |
| super().__init__() | |
| self.detector = None | |
| def load(self) -> bool: | |
| if self._loaded: | |
| return True | |
| try: | |
| from controlnet_aux import DWposeDetector | |
| self.detector = DWposeDetector.from_pretrained( | |
| "lllyasviel/ControlNet-v1-1-nightly", | |
| subfolder="annotator/ckpts" | |
| ) | |
| # DWPose loads on GPU by default if cuda available | |
| self._loaded = True | |
| return True | |
| except Exception as e: | |
| print(f"❌ Failed to load DWPose: {e}") | |
| return False | |
| def unload(self) -> None: | |
| if not self._loaded: | |
| return | |
| self.detector = None | |
| self._loaded = False | |
| import torch | |
| import gc | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| def _execute(self, inputs: dict) -> dict: | |
| image = inputs.get("image") | |
| if image is None: | |
| raise ValueError("Input 'image' is required") | |
| if self.detector is None: | |
| # Return empty skeleton | |
| empty_pose = Image.new("RGBA", image.size, (0, 0, 0, 0)) | |
| return { | |
| "pose_keypoints": [], | |
| "pose_image": empty_pose, | |
| "image": image | |
| } | |
| # Run pose detection | |
| # DWPose detector returns a PIL Image of the pose skeleton | |
| pose_img = self.detector(image) | |
| return { | |
| "pose_keypoints": [], | |
| "pose_image": pose_img, | |
| "image": image | |
| } | |