creek

  1"""
  2Subject
  3Original file: https://colab.research.google.com/drive/1izs_-Hld8X4wtGylu8y-96PJb04Nc3kz
  4"""
  5
  6!pip install -r https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt
  7!git clone https://github.com/ZhengPeng7/BiRefNet.git
  8
  9# %cd BiRefNet
 10
 11## Create .pt file
 12
 13import torch
 14from torchvision import transforms
 15from models.birefnet import BiRefNet
 16
 17device = "cuda" if torch.cuda.is_available() else "cpu"
 18print("Using device:", device)
 19
 20birefnet = BiRefNet.from_pretrained("ZhengPeng7/BiRefNet")
 21birefnet.to(device)
 22birefnet.eval()
 23birefnet.half()
 24
 25
 26model_fp16_cpu = birefnet.to("cpu")
 27torch.save(model_fp16_cpu.state_dict(), "birefnet_fp16.pt")
 28
 29print("🔥 Saved fp16 model to birefnet_fp16.pt")
 30
 31image_size = (1024, 1024)
 32transform_image = transforms.Compose(
 33    [
 34        transforms.Resize(image_size),
 35        transforms.ToTensor(),
 36        transforms.Normalize(
 37            mean=[0.485, 0.456, 0.406],
 38            std=[0.229, 0.224, 0.225],
 39        ),
 40    ]
 41)
 42
 43## Inference Code
 44
 45import os
 46import io
 47import time
 48import numpy as np
 49from PIL import Image
 50import cv2
 51import matplotlib.pyplot as plt
 52
 53import torch
 54from torchvision import transforms
 55from google.colab import files
 56
 57
 58device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 59print("Using device:", device)
 60
 61
 62OUTPUT_DIR = "/content/output/birefnet_cutouts"
 63os.makedirs(OUTPUT_DIR, exist_ok=True)
 64print("Cutouts will be saved in:", OUTPUT_DIR)
 65
 66
 67image_size = (1024, 1024)
 68transform_image = transforms.Compose(
 69    [
 70        transforms.Resize(image_size),
 71        transforms.ToTensor(),
 72        transforms.Normalize(
 73            mean=[0.485, 0.456, 0.406],
 74            std=[0.229, 0.224, 0.225],
 75        ),
 76    ]
 77)
 78
 79
 80if "birefnet" not in globals():
 81    raise RuntimeError(
 82        "Model `birefnet` not found. Load your BiRefNet model and weights first, e.g.:\n"
 83        "  birefnet = BiRefNet(...)\n"
 84        "  birefnet.load_state_dict(torch.load('ckpt.pth', map_location='cpu'))\n"
 85        "  birefnet.to(device).eval()\n"
 86    )
 87
 88birefnet = birefnet.to(device).eval()
 89
 90
 91def run_birefnet_on_image(pil_img: Image.Image, thresh: float = 0.5):
 92    """
 93    Runs BiRefNet on a single PIL image.
 94    Returns:
 95      mask_bin      : (H, W) uint8 (0 or 1) -- final model output mask
 96      prob_resized  : (H, W) float32 [0,1]  -- probability map (not saved by default)
 97      infer_time    : float seconds
 98    """
 99    orig_w, orig_h = pil_img.size
100
101    inp = transform_image(pil_img).unsqueeze(0).to(device)
102    try:
103        inp = inp.half()
104    except Exception:
105        pass
106
107    with torch.no_grad():
108        t0 = time.perf_counter()
109        preds = birefnet(inp)[-1]
110        preds = preds.float().sigmoid().cpu()
111        t1 = time.perf_counter()
112
113    prob = preds[0, 0].numpy().astype(np.float32)
114    prob_resized = cv2.resize(prob, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
115    mask_bin = (prob_resized >= thresh).astype(np.uint8)
116
117    return mask_bin, prob_resized, (t1 - t0)
118
119
120def make_binary_cutout(img_rgb: np.ndarray, mask_bin: np.ndarray):
121    """
122    img_rgb: HxWx3 uint8
123    mask_bin: HxW uint8 {0,1}
124    Returns RGBA HxWx4 uint8 with alpha=255 on foreground, 0 on background.
125    """
126    H, W = img_rgb.shape[:2]
127    if mask_bin.shape != (H, W):
128        mask_bin = cv2.resize(mask_bin, (W, H), interpolation=cv2.INTER_NEAREST)
129
130    alpha_u8 = (mask_bin * 255).astype(np.uint8)
131    rgba = np.dstack([img_rgb, alpha_u8])
132    return rgba
133
134
135print("🔼 Upload one or more images to segment with BiRefNet...")
136uploaded = files.upload()
137
138THRESH = 0.5
139KEEP_LARGEST = True
140
141for fname, file_data in uploaded.items():
142    try:
143        print(f"\n=== Processing: {fname} ===")
144        pil_img = Image.open(io.BytesIO(file_data)).convert("RGB")
145        img_np = np.array(pil_img)
146
147        mask_bin, prob_map, t_inf = run_birefnet_on_image(pil_img, thresh=THRESH)
148        print(f"BiRefNet inference time: {t_inf:.3f} s")
149
150        if KEEP_LARGEST:
151            num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
152                mask_bin, connectivity=8
153            )
154            if num_labels > 1:
155                areas = stats[1:, cv2.CC_STAT_AREA]
156                max_idx = 1 + np.argmax(areas)
157                mask_bin = (labels == max_idx).astype(np.uint8)
158
159        rgba_cutout = make_binary_cutout(img_np, mask_bin)
160
161        base = os.path.splitext(fname)[0]
162        out_path = os.path.join(OUTPUT_DIR, base + "_birefnet_cutout.png")
163        Image.fromarray(rgba_cutout).save(out_path, format="PNG")
164        print(" Saved final model cutout:", out_path)
165
166        fig, axs = plt.subplots(1, 3, figsize=(12, 4))
167        axs[0].imshow(img_np)
168        axs[0].set_title("Original")
169        axs[0].axis("off")
170        axs[1].imshow(mask_bin, cmap="gray")
171        axs[1].set_title("Final model mask")
172        axs[1].axis("off")
173        axs[2].imshow(rgba_cutout)
174        axs[2].set_title("Cutout (PNG, no background)")
175        axs[2].axis("off")
176        plt.tight_layout()
177        plt.show()
178
179        files.download(out_path)
180
181    except Exception as e:
182        print(f" Error with {fname}: {e}")
183
184print("\n Done! All cutouts saved in:", OUTPUT_DIR)