diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 30361b5d..b9ee0239 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -20,6 +20,7 @@ class OutputKeys(object): KEYPOINTS = 'keypoints' MASKS = 'masks' DEPTHS = 'depths' + DEPTHS_COLOR = 'depths_color' TEXT = 'text' POLYGONS = 'polygons' OUTPUT = 'output' diff --git a/modelscope/pipelines/cv/image_depth_estimation_pipeline.py b/modelscope/pipelines/cv/image_depth_estimation_pipeline.py index 1f580733..a5445692 100644 --- a/modelscope/pipelines/cv/image_depth_estimation_pipeline.py +++ b/modelscope/pipelines/cv/image_depth_estimation_pipeline.py @@ -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 diff --git a/tests/pipelines/test_image_depth_estimation.py b/tests/pipelines/test_image_depth_estimation.py index 933ce7a0..6ec16a64 100644 --- a/tests/pipelines/test_image_depth_estimation.py +++ b/tests/pipelines/test_image_depth_estimation.py @@ -24,9 +24,8 @@ class ImageDepthEstimationTest(unittest.TestCase, DemoCompatibilityCheck): input_location = 'data/test/images/image_depth_estimation.jpg' estimator = pipeline(Tasks.image_depth_estimation, model=self.model_id) result = estimator(input_location) - depths = result[OutputKeys.DEPTHS] - depth_viz = depth_to_color(depths) - cv2.imwrite('result.jpg', depth_viz) + depth_vis = result[OutputKeys.DEPTHS_COLOR] + cv2.imwrite('result.jpg', depth_vis) print('test_image_depth_estimation DONE')