Update README.md

README.md

@@ -124,6 +124,10 @@ pillow_image = pipe(image_path) # applies mask on input and returns a pillow ima
 
 Or load the model 
 ```python
+from PIL import Image
+from skimage import io
+import torch
+import torch.nn.functional as F
 from transformers import AutoModelForImageSegmentation
 from torchvision.transforms.functional import normalize
 model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4",trust_remote_code=True)
@@ -153,6 +157,7 @@ model.to(device)
 image_path = "https://farm5.staticflickr.com/4007/4322154488_997e69e4cf_z.jpg"
 orig_im = io.imread(image_path)
 orig_im_size = orig_im.shape[0:2]
+model_input_size = [1024, 1024]
 image = preprocess_image(orig_im, model_input_size).to(device)
 
 # inference