Upload external/MV-Adapter/mvadapter/utils/mesh_utils/cv_ops.py with huggingface_hub

external/MV-Adapter/mvadapter/utils/mesh_utils/cv_ops.py

@@ -1,9 +1,12 @@
 from typing import Optional
 
-import cv2
-import numpy as np
+import cvcuda
 import torch
 
+torch_to_cvc = lambda x, layout: cvcuda.as_tensor(x, layout)
+
+cvc_to_torch = lambda x, device: torch.tensor(x.cuda(), device=device)
+
 
 def inpaint_cvc(
     image: torch.Tensor,
@@ -14,11 +17,18 @@ def inpaint_cvc(
     input_dtype = image.dtype
     input_device = image.device
 
-    img_np = (image.detach().float().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
-    msk_np = (mask.detach().float().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+    image = image.detach()
+    mask = mask.detach()
 
-    result = cv2.inpaint(img_np, msk_np, padding_size, cv2.INPAINT_TELEA)
-    output = torch.from_numpy(result).to(device=input_device)
+    if image.dtype != torch.uint8:
+        image = (image * 255).to(torch.uint8)
+    if mask.dtype != torch.uint8:
+        mask = (mask * 255).to(torch.uint8)
+
+    image_cvc = torch_to_cvc(image, "HWC")
+    mask_cvc = torch_to_cvc(mask, "HW")
+    output_cvc = cvcuda.inpaint(image_cvc, mask_cvc, padding_size)
+    output = cvc_to_torch(output_cvc, device=input_device)
 
     if return_dtype == torch.uint8 or input_dtype == torch.uint8:
         return output
@@ -31,11 +41,14 @@ def batch_inpaint_cvc(
     padding_size: int,
     return_dtype: Optional[torch.dtype] = None,
 ):
-    return torch.stack(
-        [inpaint_cvc(image, mask, padding_size, return_dtype)
-         for image, mask in zip(images, masks)],
-        dim=0,
+    output = torch.stack(
+        [
+            inpaint_cvc(image, mask, padding_size, return_dtype)
+            for (image, mask) in zip(images, masks)
+        ],
+        axis=0,
     )
+    return output
 
 
 def batch_erode(
@@ -43,19 +56,20 @@ def batch_erode(
 ):
     input_dtype = masks.dtype
     input_device = masks.device
-
-    msk_np = (masks.detach().float().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
-
-    if msk_np.ndim == 3:
-        eroded = np.stack([cv2.erode(m, kernel) for m in msk_np])
-    else:
-        eroded = cv2.erode(msk_np, kernel)
-
-    output = torch.from_numpy(eroded).to(device=input_device)
+    masks = masks.detach()
+    if masks.dtype != torch.uint8:
+        masks = (masks.float() * 255).to(torch.uint8)
+    masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
+    masks_erode_cvc = cvcuda.morphology(
+        masks_cvc,
+        cvcuda.MorphologyType.ERODE,
+        maskSize=(kernel_size, kernel_size),
+        anchor=(-1, -1),
+    )
+    masks_erode = cvc_to_torch(masks_erode_cvc, device=input_device)[..., 0]
     if return_dtype == torch.uint8 or input_dtype == torch.uint8:
-        return output
-    return (output > 0).to(dtype=input_dtype)
+        return masks_erode
+    return (masks_erode > 0).to(dtype=input_dtype)
 
 
 def batch_dilate(
@@ -63,16 +77,17 @@ def batch_dilate(
 ):
     input_dtype = masks.dtype
     input_device = masks.device
-
-    msk_np = (masks.detach().float().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
-
-    if msk_np.ndim == 3:
-        dilated = np.stack([cv2.dilate(m, kernel) for m in msk_np])
-    else:
-        dilated = cv2.dilate(msk_np, kernel)
-
-    output = torch.from_numpy(dilated).to(device=input_device)
+    masks = masks.detach()
+    if masks.dtype != torch.uint8:
+        masks = (masks.float() * 255).to(torch.uint8)
+    masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
+    masks_dilate_cvc = cvcuda.morphology(
+        masks_cvc,
+        cvcuda.MorphologyType.DILATE,
+        maskSize=(kernel_size, kernel_size),
+        anchor=(-1, -1),
+    )
+    masks_dilate = cvc_to_torch(masks_dilate_cvc, device=input_device)[..., 0]
     if return_dtype == torch.uint8 or input_dtype == torch.uint8:
-        return output
-    return (output > 0).to(dtype=input_dtype)
+        return masks_dilate
+    return (masks_dilate > 0).to(dtype=input_dtype)