Upload patches/TripoSG_image_process.py with huggingface_hub

patches/TripoSG_image_process.py

@@ -0,0 +1,165 @@
+# -*- coding: utf-8 -*-
+# Patched: resize BEFORE squeeze to fix "ValueError: spatial dimensions of [1024]"
+# Original bug: alpha_gpu_rmbg was squeezed to [H,W] before resize_transform([H,W])
+#               which expects [C,H,W] input. Fix: resize first, then squeeze(0).
+import os
+from skimage.morphology import remove_small_objects
+from skimage.measure import label
+import numpy as np
+from PIL import Image
+import cv2
+from torchvision import transforms
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+def find_bounding_box(gray_image):
+    _, binary_image = cv2.threshold(gray_image, 1, 255, cv2.THRESH_BINARY)
+    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    max_contour = max(contours, key=cv2.contourArea)
+    x, y, w, h = cv2.boundingRect(max_contour)
+    return x, y, w, h
+
+def load_image(img_path, bg_color=None, rmbg_net=None, padding_ratio=0.1):
+    img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+    if img is None:
+        return f"invalid image path {img_path}"
+
+    def is_valid_alpha(alpha, min_ratio = 0.01):
+        bins = 20
+        if isinstance(alpha, np.ndarray):
+            hist = cv2.calcHist([alpha], [0], None, [bins], [0, 256])
+        else:
+            hist = torch.histc(alpha, bins=bins, min=0, max=1)
+        min_hist_val = alpha.shape[0] * alpha.shape[1] * min_ratio
+        return hist[0] >= min_hist_val and hist[-1] >= min_hist_val
+
+    def rmbg(image: torch.Tensor) -> torch.Tensor:
+        model_class = type(rmbg_net).__name__
+        if 'BriaRMBG' in model_class:
+            # RMBG-1.4: normalize to [-0.5, 0.5] range, take finest prediction d1
+            img_v1 = TF.normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]).unsqueeze(0)
+            result = rmbg_net(img_v1)
+            # BriaRMBG returns (d1, d2, ..., d7); d1 is finest
+            d1 = result[0] if isinstance(result, (list, tuple)) else result
+            return d1[0]  # [B,1,H,W] -> [1,H,W]
+        else:
+            # RMBG-2.0 (BiRefNet): ImageNet normalization, last item is finest
+            img_v2 = TF.normalize(image, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).unsqueeze(0)
+            result = rmbg_net(img_v2)
+            last = result[-1] if isinstance(result, (list, tuple)) else result
+            return last[0]  # [B,1,H,W] -> [1,H,W]
+
+    if len(img.shape) == 2:
+        num_channels = 1
+    else:
+        num_channels = img.shape[2]
+
+    # check if too large
+    height, width = img.shape[:2]
+    if height > width:
+        scale = 2000 / height
+    else:
+        scale = 2000 / width
+    if scale < 1:
+        new_size = (int(width * scale), int(height * scale))
+        img = cv2.resize(img, new_size, interpolation=cv2.INTER_AREA)
+
+    if img.dtype != 'uint8':
+        img = (img * (255. / np.iinfo(img.dtype).max)).astype(np.uint8)
+
+    rgb_image = None
+    alpha = None
+
+    if num_channels == 1:
+        rgb_image = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    elif num_channels == 3:
+        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    elif num_channels == 4:
+        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+
+        b, g, r, alpha = cv2.split(img)
+        if not is_valid_alpha(alpha):
+            alpha = None
+        else:
+            alpha_gpu = torch.from_numpy(alpha).unsqueeze(0).cuda().float() / 255.
+    else:
+        return f"invalid image: channels {num_channels}"
+
+    rgb_image_gpu = torch.from_numpy(rgb_image).cuda().float().permute(2, 0, 1) / 255.
+    if alpha is None:
+        resize_transform = transforms.Resize((384, 384), antialias=True)
+        rgb_image_resized = resize_transform(rgb_image_gpu)
+        normalize_image = rgb_image_resized * 2 - 1
+
+        mean_color = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).cuda()
+        resize_transform = transforms.Resize((1024, 1024), antialias=True)
+        rgb_image_resized = resize_transform(rgb_image_gpu)
+        max_value = rgb_image_resized.flatten().max()
+        if max_value < 1e-3:
+            return "invalid image: pure black image"
+        normalize_image = rgb_image_resized / max_value - mean_color
+        normalize_image = normalize_image.unsqueeze(0)
+        resize_transform = transforms.Resize((rgb_image_gpu.shape[1], rgb_image_gpu.shape[2]), antialias=True)
+
+        # seg from rmbg
+        alpha_gpu_rmbg = rmbg(rgb_image_resized)
+        # FIX: resize FIRST (needs [C,H,W]), THEN squeeze to [H,W]
+        alpha_gpu_rmbg = resize_transform(alpha_gpu_rmbg)
+        alpha_gpu_rmbg = alpha_gpu_rmbg.squeeze(0)
+        mi = alpha_gpu_rmbg.min()
+        ma = alpha_gpu_rmbg.max()
+        alpha_gpu_rmbg = (alpha_gpu_rmbg - mi) / (ma - mi)
+
+        alpha_gpu = alpha_gpu_rmbg
+
+        alpha_gpu_tmp = alpha_gpu * 255
+        alpha = alpha_gpu_tmp.to(torch.uint8).squeeze().cpu().numpy()
+
+        _, alpha = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
+        labeled_alpha = label(alpha)
+        cleaned_alpha = remove_small_objects(labeled_alpha, min_size=200)
+        cleaned_alpha = (cleaned_alpha > 0).astype(np.uint8)
+        alpha = cleaned_alpha * 255
+        alpha_gpu = torch.from_numpy(cleaned_alpha).cuda().float().unsqueeze(0)
+        x, y, w, h = find_bounding_box(alpha)
+
+    # If alpha is provided, the bounds of all foreground are used
+    else:
+        rows, cols = np.where(alpha > 0)
+        if rows.size > 0 and cols.size > 0:
+            x_min = np.min(cols)
+            y_min = np.min(rows)
+            x_max = np.max(cols)
+            y_max = np.max(rows)
+
+            width = x_max - x_min + 1
+            height = y_max - y_min + 1
+        x, y, w, h = x_min, y_min, width, height
+
+    if np.all(alpha==0):
+        raise ValueError(f"input image too small")
+
+    bg_gray = bg_color[0]
+    bg_color = torch.from_numpy(bg_color).float().cuda().repeat(alpha_gpu.shape[1], alpha_gpu.shape[2], 1).permute(2, 0, 1)
+    rgb_image_gpu = rgb_image_gpu * alpha_gpu + bg_color * (1 - alpha_gpu)
+    padding_size = [0] * 6
+    if w > h:
+        padding_size[0] = int(w * padding_ratio)
+        padding_size[2] = int(padding_size[0] + (w - h) / 2)
+    else:
+        padding_size[2] = int(h * padding_ratio)
+        padding_size[0] = int(padding_size[2] + (h - w) / 2)
+    padding_size[1] = padding_size[0]
+    padding_size[3] = padding_size[2]
+    padded_tensor = F.pad(rgb_image_gpu[:, y:(y+h), x:(x+w)], pad=tuple(padding_size), mode='constant', value=bg_gray)
+
+    return padded_tensor
+
+def prepare_image(image_path, bg_color, rmbg_net=None):
+    if os.path.isfile(image_path):
+        img_tensor = load_image(image_path, bg_color=bg_color, rmbg_net=rmbg_net)
+        img_np = img_tensor.permute(1,2,0).cpu().numpy()
+        img_pil = Image.fromarray((img_np*255).astype(np.uint8))
+
+        return img_pil