fix: Move the manual image resizing out of the depth anything pipeline

This commit is contained in:
blessedcoolant 2024-07-31 23:35:52 +05:30
parent 13fb2d1f49
commit daf899f9c4
2 changed files with 10 additions and 7 deletions

View File

@ -594,6 +594,7 @@ class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small", "small_v2"]
# DepthAnything V2 Small model is licensed under Apache 2.0 but not the base and large models.
DEPTH_ANYTHING_MODELS = {
"large": "LiheYoung/depth-anything-large-hf",
"base": "LiheYoung/depth-anything-base-hf",
@ -627,7 +628,13 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
source=DEPTH_ANYTHING_MODELS[self.model_size], loader=load_depth_anything
) as depth_anything_detector:
assert isinstance(depth_anything_detector, DepthAnythingPipeline)
return depth_anything_detector.generate_depth(image, self.resolution)
depth_map = depth_anything_detector.generate_depth(image)
# Resizing to user target specified size
new_height = int(image.size[1] * (self.resolution / image.size[0]))
depth_map = depth_map.resize((self.resolution, new_height))
return depth_map
@invocation(

View File

@ -1,4 +1,4 @@
from typing import Optional, cast
from typing import Optional
import torch
from PIL import Image
@ -14,13 +14,9 @@ class DepthAnythingPipeline(RawModel):
def __init__(self, pipeline: DepthEstimationPipeline) -> None:
self.pipeline = pipeline
def generate_depth(self, image: Image.Image, resolution: int = 512):
image_width, image_height = image.size
def generate_depth(self, image: Image.Image) -> Image.Image:
depth_map = self.pipeline(image)["depth"]
assert isinstance(depth_map, Image.Image)
new_height = int(image_height * (resolution / image_width))
depth_map = depth_map.resize((resolution, new_height))
return depth_map
def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):