Create marigold_depth_estimation.py

marigold_depth_estimation.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import DiffusionPipeline, UNet2DConditionModel, AutoencoderKL, DDIMScheduler
+from diffusers.utils import BaseOutput
+
+class MarigoldDepthOutput(BaseOutput):
+    depth_np: np.ndarray
+    depth_image: Image.Image
+
+class MarigoldPipeline(DiffusionPipeline):
+    def __init__(self, unet: UNet2DConditionModel, vae: AutoencoderKL, scheduler: DDIMScheduler):
+        super().__init__()
+        self.unet = unet
+        self.vae = vae
+        self.scheduler = scheduler
+
+    @torch.no_grad()
+    def __call__(self, input_image: Image, denoising_steps: int = 10, save_path: str = None) -> MarigoldDepthOutput:
+        device = self.device
+
+        # Image preprocessing
+        input_image = input_image.convert("RGB")
+        image = np.asarray(input_image)
+        rgb = np.transpose(image, (2, 0, 1))
+        rgb_norm = rgb / 255.0 * 2.0 - 1.0
+        rgb_norm = torch.from_numpy(rgb_norm).to(device)
+
+        # Encode image
+        rgb_latent = self._encode_rgb(rgb_norm)
+
+        # Initial depth map (noise)
+        depth_latent = torch.randn(rgb_latent.shape, device=device)
+
+        # Denoising loop
+        timesteps = self.scheduler.timesteps
+        for t in timesteps:
+            unet_input = torch.cat([rgb_latent, depth_latent], dim=1)
+            noise_pred = self.unet(unet_input, t).sample
+            depth_latent = self.scheduler.step(noise_pred, t, depth_latent).prev_sample
+
+        # Decode depth map
+        depth = self._decode_depth(depth_latent)
+
+        # Scale to [0, 1] and convert to numpy
+        depth = (depth + 1.0) / 2.0
+        depth_np = depth.cpu().numpy().astype(np.float32)
+        depth_image = (depth_np * 255).astype(np.uint8)
+        depth_image = Image.fromarray(depth_image[0], 'L')  # 'L' mode for grayscale image
+
+        # Save the depth map image if a path is provided
+        if save_path:
+            depth_image.save(save_path)
+
+        return MarigoldDepthOutput(depth_np=depth_np, depth_image=depth_image)
+
+    def _encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
+        h = self.vae.encoder(rgb_in)
+        moments = self.vae.quant_conv(h)
+        mean, _ = torch.chunk(moments, 2, dim=1)
+        rgb_latent = mean * 0.18215
+        return rgb_latent
+
+    def _decode_depth(self, depth_latent: torch.Tensor) -> torch.Tensor:
+        z = self.vae.post_quant_conv(depth_latent)
+        stacked = self.vae.decoder(z)
+        depth_mean = stacked.mean(dim=1, keepdim=True)
+        return depth_mean
+
+# Instantiate the model components and the pipeline
+unet_model = UNet2DConditionModel()
+vae_model = AutoencoderKL()
+scheduler = DDIMScheduler()
+pipeline = MarigoldPipeline(unet=unet_model, vae=vae_model, scheduler=scheduler)
+
+# Load an image and predict the depth map
+input_image = Image.open('path_to_your_image.jpg')
+output_path = 'path_to_save_image.jpg'  # Specify the path where you want to save the depth image
+output = pipeline(input_image, denoising_steps=10, save_path=output_path)