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

@ -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):