1 # Load BiRefNet with weights
2 from transformers import AutoModelForImageSegmentation
3 birefnet = AutoModelForImageSegmentation . from_pretrained ( 'ZhengPeng7/BiRefNet_HR' , trust_remote_code = True )
1 # Download codes
2 git clone https://github.com/ZhengPeng7/BiRefNet.git
3 cd BiRefNet
1 # Use codes locally
2 from models . birefnet import BiRefNet
3
4 # Load weights from Hugging Face Models
5 birefnet = BiRefNet . from_pretrained ( 'ZhengPeng7/BiRefNet_HR' )
1 # Use codes and weights locally
2 import torch
3 from utils import check_state_dict
4
5 birefnet = BiRefNet ( bb_pretrained = False )
6 state_dict = torch . load ( PATH_TO_WEIGHT , map_location = 'cpu' )
7 state_dict = check_state_dict ( state_dict )
8 birefnet . load_state_dict ( state_dict )
1 # Imports
2 from PIL import Image
3 import matplotlib . pyplot as plt
4 import torch
5 from torchvision import transforms
6 from models . birefnet import BiRefNet
7
8 birefnet = . . . # -- BiRefNet should be loaded with codes above, either way.
9 torch . set_float32_matmul_precision ( [ 'high' , 'highest' ] [ 0 ] )
10 birefnet . to ( 'cuda' )
11 birefnet . eval ( )
12 birefnet . half ( )
13
14 def extract_object ( birefnet , imagepath ) :
15 # Data settings
16 image_size = ( 2048 , 2048 )
17 transform_image = transforms . Compose ( [
18 transforms . Resize ( image_size ) ,
19 transforms . ToTensor ( ) ,
20 transforms . Normalize ( [ 0.485 , 0.456 , 0.406 ] , [ 0.229 , 0.224 , 0.225 ] )
21 ] )
22
23 image = Image . open ( imagepath )
24 input_images = transform_image ( image ) . unsqueeze ( 0 ) . to ( 'cuda' ) . half ( )
25
26 # Prediction
27 with torch . no_grad ( ) :
28 preds = birefnet ( input_images ) [ - 1 ] . sigmoid ( ) . cpu ( )
29 pred = preds [ 0 ] . squeeze ( )
30 pred_pil = transforms . ToPILImage ( ) ( pred )
31 mask = pred_pil . resize ( image . size )
32 image . putalpha ( mask )
33 return image , mask
34
35 # Visualization
36 plt . axis ( "off" )
37 plt . imshow ( extract_object ( birefnet , imagepath = 'PATH-TO-YOUR_IMAGE.jpg' ) [ 0 ] )
38 plt . show ( )
39
import requests
import base64
from io import BytesIO
from PIL import Image
YOUR_HF_TOKEN = 'xxx'
API_URL = "xxx"
headers = {
"Authorization": "Bearer {}".format(YOUR_HF_TOKEN)
}
def base64_to_bytes(base64_string):
# Remove the data URI prefix if present
if "data:image" in base64_string:
base64_string = base64_string.split(",")[1]
# Decode the Base64 string into bytes
image_bytes = base64.b64decode(base64_string)
return image_bytes
def bytes_to_base64(image_bytes):
# Create a BytesIO object to handle the image data
image_stream = BytesIO(image_bytes)
# Open the image using Pillow (PIL)
image = Image.open(image_stream)
return image
def query(payload):
response = requests.post(API_URL, headers=headers, json=payload)
return response.json()
output = query({
"inputs": "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg",
"parameters": {}
})
output_image = bytes_to_base64(base64_to_bytes(output))
output_imageThis repo contains the weights of BiRefNet proposed in our paper, which has achieved the SOTA performance on three tasks (DIS, HRSOD, and COD).
Go to my GitHub page for BiRefNet codes and the latest updates:
https://github.com/ZhengPeng7/BiRefNet :)
@article{zheng2024birefnet,
title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
journal={CAAI Artificial Intelligence Research},
volume = {3},
pages = {9150038},
year={2024}
}
