Update README.md

README.md

@@ -40,7 +40,7 @@ BEN2 was trained on the DIS5k and our 22K proprietary segmentation dataset. Our
 ## Quick start code (inside cloned repo)
 
 ```python
-import model
+import BEN2
 from PIL import Image
 import torch
 
@@ -49,7 +49,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 file = "./image.png" # input image
 
-model = model.BEN_Base().to(device).eval() #init pipeline
+model = BEN2.BEN_Base().to(device).eval() #init pipeline
 
 model.loadcheckpoints("./BEN2_Base.pth")
 image = Image.open(file)
@@ -63,7 +63,7 @@ foreground.save("./foreground.png")
 ## Batch image processing
 
 ```python
-import model
+import BEN2
 from PIL import Image
 import torch
 
@@ -72,7 +72,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 
-model = model.BEN_Base().to(device).eval() #init pipeline
+model = BEN2.BEN_Base().to(device).eval() #init pipeline
 
 model.loadcheckpoints("./BEN2_Base.pth")
 
@@ -102,7 +102,7 @@ sudo apt install ffmpeg
 ```
 
 ```python
-import model
+import BEN2
 from PIL import Image
 import torch
 
@@ -111,7 +111,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 file = "./image.png" # input image
 
-model = model.BEN_Base().to(device).eval() #init pipeline
+model = BEN2.BEN_Base().to(device).eval() #init pipeline
 
 model.loadcheckpoints("./BEN2_Base.pth")