Refactor DepthPro object

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

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)

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__":

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)