90 lines
2.4 KiB
Python
90 lines
2.4 KiB
Python
# Imports
|
|
import pdb
|
|
import time
|
|
|
|
import torch
|
|
import tqlt.utils as tu
|
|
from models.birefnet import BiRefNet
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
# # Option 1: loading BiRefNet with weights:
|
|
from transformers import AutoModelForImageSegmentation
|
|
|
|
# # Option-3: Loading model and weights from local disk:
|
|
from utils import check_state_dict
|
|
|
|
# birefnet = AutoModelForImageSegmentation.from_pretrained(
|
|
# "zhengpeng7/BiRefNet", trust_remote_code=True, local
|
|
# )
|
|
|
|
# # Option-2: loading weights with BiReNet codes:
|
|
# birefnet = BiRefNet.from_pretrained('zhengpeng7/BiRefNet')
|
|
imgs = tu.next_files("./in_the_wild", ".png")
|
|
|
|
|
|
birefnet = BiRefNet(bb_pretrained=False)
|
|
state_dict = torch.load("./BiRefNet-general-epoch_244.pth", map_location="cpu")
|
|
state_dict = check_state_dict(state_dict)
|
|
birefnet.load_state_dict(state_dict)
|
|
|
|
|
|
# Load Model
|
|
device = "cuda"
|
|
torch.set_float32_matmul_precision(["high", "highest"][0])
|
|
|
|
birefnet.to(device)
|
|
birefnet.eval()
|
|
print("BiRefNet is ready to use.")
|
|
|
|
# Input Data
|
|
transform_image = transforms.Compose(
|
|
[
|
|
transforms.Resize((1024, 1024)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
|
]
|
|
)
|
|
|
|
|
|
import os
|
|
from glob import glob
|
|
|
|
from image_proc import refine_foreground
|
|
|
|
src_dir = "./images_todo"
|
|
image_paths = glob(os.path.join(src_dir, "*"))
|
|
dst_dir = "./predictions"
|
|
os.makedirs(dst_dir, exist_ok=True)
|
|
for image_path in imgs:
|
|
|
|
print("Processing {} ...".format(image_path))
|
|
image = Image.open(image_path)
|
|
input_images = transform_image(image).unsqueeze(0).to("cuda")
|
|
|
|
# Prediction
|
|
start = time.time()
|
|
|
|
with torch.no_grad():
|
|
preds = birefnet(input_images)[-1].sigmoid().cpu()
|
|
|
|
print(time.time() - start)
|
|
pred = preds[0].squeeze()
|
|
|
|
# Save Results
|
|
file_ext = os.path.splitext(image_path)[-1]
|
|
pred_pil = transforms.ToPILImage()(pred)
|
|
pred_pil = pred_pil.resize(image.size)
|
|
pred_pil.save(image_path.replace(src_dir, dst_dir).replace(file_ext, "-mask.png"))
|
|
image_masked = refine_foreground(image, pred_pil)
|
|
image_masked.putalpha(pred_pil)
|
|
image_masked.save(
|
|
image_path.replace(src_dir, dst_dir).replace(file_ext, "-subject.png")
|
|
)
|
|
|
|
# Save Results
|
|
file_ext = os.path.splitext(image_path)[-1]
|
|
pred_pil = transforms.ToPILImage()(pred)
|
|
pred_pil = pred_pil.resize(image.size)
|
|
pred_pil.save(image_path.replace(src_dir, dst_dir).replace(file_ext, "-mask.png"))
|