69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
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
|
|
|
|
from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
|
|
from mflux.models.depth_pro.depth_pro_model import DepthProModel
|
|
from mflux.post_processing.image_util import ImageUtil
|
|
|
|
|
|
@dataclass
|
|
class DepthResult:
|
|
depth_image: Image.Image
|
|
depth_array: mx.array
|
|
min_depth: float
|
|
max_depth: float
|
|
|
|
|
|
class DepthPro(nn.Module):
|
|
def __init__(self, quantize: int | None = None):
|
|
super().__init__()
|
|
self.depth_pro_model = DepthProModel()
|
|
DepthProInitializer.init(self.depth_pro_model, quantize=quantize)
|
|
|
|
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)
|
|
|
|
@staticmethod
|
|
def _pre_process(image_path):
|
|
image = Image.open(image_path).convert("RGB")
|
|
input_array = ImageUtil.preprocess_for_depth_pro(image)
|
|
input_array = DepthPro._resize(input_array)
|
|
return input_array, image.height, image.width
|
|
|
|
@staticmethod
|
|
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)
|
|
depth_np = np.asarray(normalized_depth.squeeze())
|
|
depth_image = Image.fromarray((depth_np * 255).astype(np.uint8))
|
|
depth_image = depth_image.resize((width, height))
|
|
depth_np = np.array(depth_image) / 255.0
|
|
return DepthResult(
|
|
depth_image=depth_image,
|
|
depth_array=depth,
|
|
min_depth=depth_min.item(),
|
|
max_depth=depth_max.item(),
|
|
)
|
|
|
|
@staticmethod
|
|
def _resize(x: mx.array) -> mx.array:
|
|
x_np = np.array(x)
|
|
x_torch = torch.from_numpy(x_np)
|
|
x_torch = x_torch.unsqueeze(0)
|
|
x_torch = torch.nn.functional.interpolate(
|
|
x_torch,
|
|
size=(1536, 1536),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return mx.array(x_torch)
|