ZhengPeng7
/

BiRefNet_lite

Image Segmentation

background-removal

mask-generation

Dichotomous Image Segmentation

Camouflaged Object Detection

Salient Object Detection

pytorch_model_hub_mixin

model_hub_mixin

Model card Files Files and versions

ZhengPeng7 commited on Feb 5

Commit

7d10348

·

1 Parent(s): b9a7bec

Add half precision setting in README.

Files changed (1) hide show

README.md +2 -1

README.md CHANGED Viewed

@@ -112,6 +112,7 @@ birefnet = ... # -- BiRefNet should be loaded with codes above, either way.
 torch.set_float32_matmul_precision(['high', 'highest'][0])
 birefnet.to('cuda')
 birefnet.eval()
 def extract_object(birefnet, imagepath):
     # Data settings
@@ -123,7 +124,7 @@ def extract_object(birefnet, imagepath):
     ])
     image = Image.open(imagepath)
-    input_images = transform_image(image).unsqueeze(0).to('cuda')
     # Prediction
     with torch.no_grad():

 torch.set_float32_matmul_precision(['high', 'highest'][0])
 birefnet.to('cuda')
 birefnet.eval()
+birefnet.half()
 def extract_object(birefnet, imagepath):
     # Data settings
     ])
     image = Image.open(imagepath)
+    input_images = transform_image(image).unsqueeze(0).to('cuda').half()
     # Prediction
     with torch.no_grad():