[ldm3d] Update code to be functional with the new checkpoints (huggingface#3875)

* fixed typo

* updated doc to be consistent in naming

* make style/quality

* preprocessing for 4 channels and not 6

* make style

* test for 4c

* make style/quality

* fixed test on cpu

---------

Co-authored-by: Aflalo <[email protected]>
Co-authored-by: Aflalo <[email protected]>
Co-authored-by: Aflalo <[email protected]>

Author: 4 people

image_processor.py

@@ -312,12 +312,17 @@ def numpy_to_depth(self, images):
         """
         if images.ndim == 3:
             images = images[None, ...]
-        images = (images * 255).round().astype("uint8")
-        if images.shape[-1] == 1:
-            # special case for grayscale (single channel) images
-            raise Exception("Not supported")
+        images_depth = images[:, :, :, 3:]
+        if images.shape[-1] == 6:
+            images_depth = (images_depth * 255).round().astype("uint8")
+            pil_images = [
+                Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
+            ]
+        elif images.shape[-1] == 4:
+            images_depth = (images_depth * 65535.0).astype(np.uint16)
+            pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
         else:
-            pil_images = [Image.fromarray(self.rgblike_to_depthmap(image[:, :, 3:]), mode="I;16") for image in images]
+            raise Exception("Not supported")
 
         return pil_images
 
@@ -349,7 +354,11 @@ def postprocess(
         image = self.pt_to_numpy(image)
 
         if output_type == "np":
-            return image[:, :, :, :3], np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
+            if image.shape[-1] == 6:
+                image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
+            else:
+                image_depth = image[:, :, :, 3:]
+            return image[:, :, :, :3], image_depth
 
         if output_type == "pil":
             return self.numpy_to_pil(image), self.numpy_to_depth(image)