creek

"""
Subject
Original file: https://colab.research.google.com/drive/1izs_-Hld8X4wtGylu8y-96PJb04Nc3kz
"""

!pip install -r https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt
!git clone https://github.com/ZhengPeng7/BiRefNet.git

# %cd BiRefNet

## Create .pt file

import torch
from torchvision import transforms
from models.birefnet import BiRefNet

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

birefnet = BiRefNet.from_pretrained("ZhengPeng7/BiRefNet")
birefnet.to(device)
birefnet.eval()
birefnet.half()


model_fp16_cpu = birefnet.to("cpu")
torch.save(model_fp16_cpu.state_dict(), "birefnet_fp16.pt")

print("🔥 Saved fp16 model to birefnet_fp16.pt")

image_size = (1024, 1024)
transform_image = transforms.Compose(
    [
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ]
)

## Inference Code

import os
import io
import time
import numpy as np
from PIL import Image
import cv2
import matplotlib.pyplot as plt

import torch
from torchvision import transforms
from google.colab import files


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)


OUTPUT_DIR = "/content/output/birefnet_cutouts"
os.makedirs(OUTPUT_DIR, exist_ok=True)
print("Cutouts will be saved in:", OUTPUT_DIR)


image_size = (1024, 1024)
transform_image = transforms.Compose(
    [
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ]
)


if "birefnet" not in globals():
    raise RuntimeError(
        "Model `birefnet` not found. Load your BiRefNet model and weights first, e.g.:\n"
        "  birefnet = BiRefNet(...)\n"
        "  birefnet.load_state_dict(torch.load('ckpt.pth', map_location='cpu'))\n"
        "  birefnet.to(device).eval()\n"
    )

birefnet = birefnet.to(device).eval()


def run_birefnet_on_image(pil_img: Image.Image, thresh: float = 0.5):
    """
    Runs BiRefNet on a single PIL image.
    Returns:
      mask_bin      : (H, W) uint8 (0 or 1) -- final model output mask
      prob_resized  : (H, W) float32 [0,1]  -- probability map (not saved by default)
      infer_time    : float seconds
    """
    orig_w, orig_h = pil_img.size

    inp = transform_image(pil_img).unsqueeze(0).to(device)
    try:
        inp = inp.half()
    except Exception:
        pass

    with torch.no_grad():
        t0 = time.perf_counter()
        preds = birefnet(inp)[-1]
        preds = preds.float().sigmoid().cpu()
        t1 = time.perf_counter()

    prob = preds[0, 0].numpy().astype(np.float32)
    prob_resized = cv2.resize(prob, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
    mask_bin = (prob_resized >= thresh).astype(np.uint8)

    return mask_bin, prob_resized, (t1 - t0)


def make_binary_cutout(img_rgb: np.ndarray, mask_bin: np.ndarray):
    """
    img_rgb: HxWx3 uint8
    mask_bin: HxW uint8 {0,1}
    Returns RGBA HxWx4 uint8 with alpha=255 on foreground, 0 on background.
    """
    H, W = img_rgb.shape[:2]
    if mask_bin.shape != (H, W):
        mask_bin = cv2.resize(mask_bin, (W, H), interpolation=cv2.INTER_NEAREST)

    alpha_u8 = (mask_bin * 255).astype(np.uint8)
    rgba = np.dstack([img_rgb, alpha_u8])
    return rgba


print("🔼 Upload one or more images to segment with BiRefNet...")
uploaded = files.upload()

THRESH = 0.5
KEEP_LARGEST = True

for fname, file_data in uploaded.items():
    try:
        print(f"\n=== Processing: {fname} ===")
        pil_img = Image.open(io.BytesIO(file_data)).convert("RGB")
        img_np = np.array(pil_img)

        mask_bin, prob_map, t_inf = run_birefnet_on_image(pil_img, thresh=THRESH)
        print(f"BiRefNet inference time: {t_inf:.3f} s")

        if KEEP_LARGEST:
            num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
                mask_bin, connectivity=8
            )
            if num_labels > 1:
                areas = stats[1:, cv2.CC_STAT_AREA]
                max_idx = 1 + np.argmax(areas)
                mask_bin = (labels == max_idx).astype(np.uint8)

        rgba_cutout = make_binary_cutout(img_np, mask_bin)

        base = os.path.splitext(fname)[0]
        out_path = os.path.join(OUTPUT_DIR, base + "_birefnet_cutout.png")
        Image.fromarray(rgba_cutout).save(out_path, format="PNG")
        print(" Saved final model cutout:", out_path)

        fig, axs = plt.subplots(1, 3, figsize=(12, 4))
        axs[0].imshow(img_np)
        axs[0].set_title("Original")
        axs[0].axis("off")
        axs[1].imshow(mask_bin, cmap="gray")
        axs[1].set_title("Final model mask")
        axs[1].axis("off")
        axs[2].imshow(rgba_cutout)
        axs[2].set_title("Cutout (PNG, no background)")
        axs[2].axis("off")
        plt.tight_layout()
        plt.show()

        files.download(out_path)

    except Exception as e:
        print(f" Error with {fname}: {e}")

print("\n Done! All cutouts saved in:", OUTPUT_DIR)