Spaces:
Sleeping
Sleeping
File size: 1,989 Bytes
c3649b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from PIL import Image
from kaggle_gpu_server.engine.plugin_base import ModelPlugin, PluginCapability, ModelCategory
class DepthAnythingV2Plugin(ModelPlugin):
name = "depth_anything_v2"
model_id = "depth-anything/Depth-Anything-V2-Small-hf"
capability = PluginCapability.DEPTH_ESTIMATION
category = ModelCategory.LIGHTWEIGHT
vram_estimate_mb = 400
version = "1.0.0"
description = "Estimates depth using Depth-Anything-V2-Small"
def __init__(self):
super().__init__()
self.pipe = None
def load(self) -> bool:
if self._loaded:
return True
try:
from transformers import pipeline
self.pipe = pipeline("depth-estimation", model=self.model_id, device=self._device)
self._loaded = True
return True
except Exception as e:
print(f"❌ Failed to load Depth Anything: {e}")
return False
def unload(self) -> None:
if not self._loaded:
return
self.pipe = 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.pipe is None:
# Fallback empty depth map
fallback = Image.new("L", image.size, 128)
return {"depth_map": fallback, "image": image}
result = self.pipe(image)
depth_map = result["depth"]
# Ensure depth map is resized to original image dimensions if transformers pipeline modified it
if depth_map.size != image.size:
depth_map = depth_map.resize(image.size, Image.Resampling.BILINEAR)
return {
"depth_map": depth_map,
"image": image
}
|