nielsr HF staff commited on
Commit
bcd125f
1 Parent(s): 5f3a1d8

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -3
README.md CHANGED
@@ -27,17 +27,17 @@ You can use the raw model for encoding document images into a vector space, but
27
  Here is how to use this model in PyTorch:
28
 
29
  ```python
30
- from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling
31
  import torch
32
  from PIL import Image
33
 
34
  image = Image.open('path_to_your_document_image').convert('RGB')
35
 
36
- feature_extractor = BeitFeatureExtractor.from_pretrained("microsoft/dit-base")
37
  model = BeitForMaskedImageModeling.from_pretrained("microsoft/dit-base")
38
 
39
  num_patches = (model.config.image_size // model.config.patch_size) ** 2
40
- pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
41
  # create random boolean mask of shape (batch_size, num_patches)
42
  bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
43
 
27
  Here is how to use this model in PyTorch:
28
 
29
  ```python
30
+ from transformers import BeitImageProcessor, BeitForMaskedImageModeling
31
  import torch
32
  from PIL import Image
33
 
34
  image = Image.open('path_to_your_document_image').convert('RGB')
35
 
36
+ processor = BeitImageProcessor.from_pretrained("microsoft/dit-base")
37
  model = BeitForMaskedImageModeling.from_pretrained("microsoft/dit-base")
38
 
39
  num_patches = (model.config.image_size // model.config.patch_size) ** 2
40
+ pixel_values = processor(images=image, return_tensors="pt").pixel_values
41
  # create random boolean mask of shape (batch_size, num_patches)
42
  bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
43