|
|
@@ -12,6 +12,7 @@ from modelscope.pipelines.base import Input, Model, Pipeline |
|
|
|
from modelscope.pipelines.builder import PIPELINES |
|
|
|
from modelscope.preprocessors import LoadImage |
|
|
|
from modelscope.utils.constant import Tasks |
|
|
|
from modelscope.utils.cv.image_utils import depth_to_color |
|
|
|
from modelscope.utils.logger import get_logger |
|
|
|
|
|
|
|
logger = get_logger() |
|
|
@@ -50,6 +51,10 @@ class ImageDepthEstimationPipeline(Pipeline): |
|
|
|
depths = results[OutputKeys.DEPTHS] |
|
|
|
if isinstance(depths, torch.Tensor): |
|
|
|
depths = depths.detach().cpu().squeeze().numpy() |
|
|
|
outputs = {OutputKeys.DEPTHS: depths} |
|
|
|
depths_color = depth_to_color(depths) |
|
|
|
outputs = { |
|
|
|
OutputKeys.DEPTHS: depths, |
|
|
|
OutputKeys.DEPTHS_COLOR: depths_color |
|
|
|
} |
|
|
|
|
|
|
|
return outputs |