fix sparse_depth is None usecase

marigold_dc.py

@@ -37,7 +37,7 @@ class MarigoldDepthCompletionPipeline(MarigoldDepthPipeline):
         # Check inputs.
         if num_inference_steps is None:
             raise ValueError("Invalid num_inference_steps")
-        if type(sparse_depth) is not np.ndarray or sparse_depth.ndim != 2:
+        if sparse_depth is not None and (type(sparse_depth) is not np.ndarray or sparse_depth.ndim != 2):
             raise ValueError(
                 "Sparse depth should be a 2D numpy ndarray with zeros at missing positions"
             )
@@ -64,7 +64,7 @@ class MarigoldDepthCompletionPipeline(MarigoldDepthPipeline):
             dtype=self.dtype,
         )  # [N,3,PPH,PPW]
 
-        if sparse_depth.shape != original_resolution:
+        if sparse_depth is not None and sparse_depth.shape != original_resolution:
             raise ValueError(
                 f"Sparse depth dimensions ({sparse_depth.shape}) must match that of the image ({image.shape[-2:]})"
             )
@@ -76,12 +76,16 @@ class MarigoldDepthCompletionPipeline(MarigoldDepthPipeline):
         del image
 
         # Preprocess sparse depth
-        sparse_depth = torch.from_numpy(sparse_depth)[None, None].float()
-        sparse_depth = sparse_depth.to(device)
-        sparse_mask = sparse_depth > 0
-        sparse_depth = sparse_depth[sparse_mask]
-        sparse_depth_min = sparse_depth.min() if sparse_depth.numel() > 0 else 0
-        sparse_depth_max = sparse_depth.max() if sparse_depth.numel() > 0 else 1
+        if sparse_depth is not None:
+            sparse_depth = torch.from_numpy(sparse_depth)[None, None].float()
+            sparse_depth = sparse_depth.to(device)
+            sparse_mask = sparse_depth > 0
+            sparse_depth = sparse_depth[sparse_mask]
+            sparse_depth_min = sparse_depth.min() if sparse_depth.numel() > 0 else 0
+            sparse_depth_max = sparse_depth.max() if sparse_depth.numel() > 0 else 1
+        else:
+            sparse_depth_min = 0
+            sparse_depth_max = 1
 
         # Set up optimization targets
         pred_latent = torch.nn.Parameter(pred_latent, requires_grad=True)
@@ -159,7 +163,7 @@ class MarigoldDepthCompletionPipeline(MarigoldDepthPipeline):
             pred_original_sample = step_output.pred_original_sample
             current_metric_estimate = latent_to_metric(pred_original_sample)
 
-            if sparse_depth.numel() > 0:
+            if sparse_depth is not None and sparse_depth.numel() > 0:
                 loss, rmse = loss_l1l2(current_metric_estimate[sparse_mask], sparse_depth)
                 loss.backward()