From 42b4fd6d90c8bbbc7634556e5e727e5ea710d788 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 10 May 2025 13:33:02 +0200 Subject: [PATCH] Refactor DepthPro object --- src/mflux/flux_tools/depth/depth_util.py | 12 ++++++------ src/mflux/models/depth_pro/depth_pro.py | 20 +++++++++++--------- src/mflux/save_depth.py | 13 ++++++++++--- tests/depth/test_depth_pro.py | 2 +- 4 files changed, 28 insertions(+), 19 deletions(-) diff --git a/src/mflux/flux_tools/depth/depth_util.py b/src/mflux/flux_tools/depth/depth_util.py index 521122d..89fa96c 100644 --- a/src/mflux/flux_tools/depth/depth_util.py +++ b/src/mflux/flux_tools/depth/depth_util.py @@ -23,7 +23,7 @@ class DepthUtil: image_path: str | Path | None = None, depth_image_path: str | Path | None = None, ) -> tuple[mx.array, PIL.Image.Image]: - # 1. Create the depth map or use existing one + # 1. Get an existing depth map or create a new one depth_image_path, depth_image = DepthUtil.get_or_create_depth_map( depth_pro=depth_pro, image_path=image_path, @@ -47,21 +47,21 @@ class DepthUtil: depth_pro: DepthPro, image_path: str | Path | None = None, depth_map_path: str | Path | None = None, - ) -> tuple[str | Path, PIL.Image.Image | None]: + ) -> tuple[str | Path, PIL.Image.Image]: # 1. If a depth map path is provided, use it directly if depth_map_path: if not os.path.exists(depth_map_path): raise FileNotFoundError(f"Depth map file not found: {depth_map_path}") return depth_map_path, None + if not image_path: raise ValueError("Either depth_map_path or image_path must be provided") - if not os.path.exists(image_path): - raise FileNotFoundError(f"Image file not found: {image_path}") - # 2. Generate a depth map from the image using the DepthPro model - depth_result = depth_pro(image_path) + # 2. Generate a depth map from the image + depth_result = depth_pro.create_depth_map(image_path=image_path) # 3. Save the depth map to a file with the same name + _depth suffix generated_depth_path = str(image_path).rsplit(".", 1)[0] + "_depth.png" depth_result.depth_image.save(generated_depth_path) + return generated_depth_path, depth_result.depth_image diff --git a/src/mflux/models/depth_pro/depth_pro.py b/src/mflux/models/depth_pro/depth_pro.py index 4855785..d8eae0c 100644 --- a/src/mflux/models/depth_pro/depth_pro.py +++ b/src/mflux/models/depth_pro/depth_pro.py @@ -1,8 +1,8 @@ +import os from dataclasses import dataclass from pathlib import Path import mlx.core as mx -import mlx.nn as nn import numpy as np import torch from PIL import Image @@ -20,16 +20,18 @@ class DepthResult: max_depth: float -class DepthPro(nn.Module): +class DepthPro: def __init__(self, quantize: int | None = None): - super().__init__() - self.depth_pro_model = DepthProModel() - DepthProInitializer.init(self.depth_pro_model, quantize=quantize) + self._depth_pro_model = DepthProModel() + DepthProInitializer.init(self._depth_pro_model, quantize=quantize) + + def create_depth_map(self, image_path: str | Path) -> DepthResult: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") - def __call__(self, image_path: str | Path, resize: bool = True) -> DepthResult: input_array, height, width = self._pre_process(image_path) - depth = self.depth_pro_model(input_array) - return self.post_process(depth, height=height, width=width) + depth = self._depth_pro_model(input_array) + return self._post_process(depth, height=height, width=width) @staticmethod def _pre_process(image_path): @@ -39,7 +41,7 @@ class DepthPro(nn.Module): return input_array, image.height, image.width @staticmethod - def post_process(depth: mx.array, height: int, width: int): + def _post_process(depth: mx.array, height: int, width: int): depth_min = mx.min(depth) depth_max = mx.max(depth) normalized_depth = (depth - depth_min) / (depth_max - depth_min) diff --git a/src/mflux/save_depth.py b/src/mflux/save_depth.py index d38d7e2..cd54517 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/save_depth.py @@ -1,4 +1,5 @@ -from mflux.flux_tools.depth.depth_util import DepthUtil +from pathlib import Path + from mflux.models.depth_pro.depth_pro import DepthPro from mflux.ui.cli.parsers import CommandLineParser @@ -9,9 +10,15 @@ def main(): parser.add_save_depth_arguments() args = parser.parse_args() - # 1. Create and save the depth map + # 1. Create the depth map depth_pro = DepthPro(quantize=args.quantize) - DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path) + depth_result = depth_pro.create_depth_map(image_path=args.image_path) + + # 2. Save the depth map with the same name + _depth suffix + image_path = Path(args.image_path) + output_path = image_path.with_stem(f"{image_path.stem}_depth").with_suffix(".png") + depth_result.depth_image.save(output_path) + print(f"Depth map saved to: {output_path}") if __name__ == "__main__": diff --git a/tests/depth/test_depth_pro.py b/tests/depth/test_depth_pro.py index 25b2b96..17e9be2 100644 --- a/tests/depth/test_depth_pro.py +++ b/tests/depth/test_depth_pro.py @@ -20,7 +20,7 @@ class TestDepthPro: depth_pro = DepthPro() # Process the image to generate a depth map - depth_result = depth_pro(str(input_image_path)) + depth_result = depth_pro.create_depth_map(str(input_image_path)) # Save the output image depth_result.depth_image.save(output_image_path)