Add guideline to use BiRefNet. Remove codes of model.

README.md

@@ -31,9 +31,12 @@ import matplotlib.pyplot as plt
 import torch
 from torchvision import transforms
 
+from models.birefnet import BiRefNet
+
+
 # Input Data
 transform_image = transforms.Compose([
-    transforms.Resize((256, 256)),
+    transforms.Resize((1024, 1024)),
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
@@ -42,7 +45,7 @@ image = Image.open(imagepath)
 input_images = transform_image(image).unsqueeze(0).to('cuda')
 
 # Load Model
-device = '0'
+device = 'cuda'
 torch.set_float32_matmul_precision(['high', 'highest'][0])
 model = BiRefNet.from_pretrained('zhengpeng7/birefnet')
 model.to(device)
@@ -55,7 +58,7 @@ with torch.no_grad():
 pred = preds[0].squeeze()
 
 # Show Results
-plt.imshow(pred, cmap='gray')
+plt.imshow(transforms.ToPILImage()(pred).resize(image.size), cmap='gray')
 plt.show()
 
 ```