pierreguillou commited on
Commit
f49e93c
1 Parent(s): 824285e

Update files/functions.py

Browse files
Files changed (1) hide show
  1. files/functions.py +22 -2
files/functions.py CHANGED
@@ -50,6 +50,26 @@ print(os.popen(f'cat /etc/issue').read())
50
  print(os.popen(f'apt search tesseract').read())
51
  import pytesseract
52
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
53
  ## Key parameters
54
 
55
  # categories colors
@@ -231,8 +251,8 @@ def set_image_dpi_resize(image):
231
  length_x, width_y = image.size
232
  factor = min(1, float(1024.0 / length_x))
233
  size = int(factor * length_x), int(factor * width_y)
234
- image_resize = image.resize(size, Image.Resampling.LANCZOS)
235
- # image_resize = image.resize(size, Image.LANCZOS)
236
  temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='1.png')
237
  temp_filename = temp_file.name
238
  image_resize.save(temp_filename, dpi=(300, 300))
 
50
  print(os.popen(f'apt search tesseract').read())
51
  import pytesseract
52
 
53
+ ## model / feature extractor / tokenizer
54
+
55
+ from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast,
56
+
57
+ import torch
58
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
59
+
60
+ # model
61
+ # tokenizer = LayoutXLMTokenizerFast.from_pretrained(model_id)
62
+ model = LayoutLMv2ForTokenClassification.from_pretrained(model_id);
63
+ model.to(device);
64
+
65
+ # feature extractor
66
+ from transformers import LayoutLMv2FeatureExtractor
67
+ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
68
+
69
+ # tokenizer
70
+ from transformers import AutoTokenizer
71
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
72
+
73
  ## Key parameters
74
 
75
  # categories colors
 
251
  length_x, width_y = image.size
252
  factor = min(1, float(1024.0 / length_x))
253
  size = int(factor * length_x), int(factor * width_y)
254
+ # image_resize = image.resize(size, Image.Resampling.LANCZOS)
255
+ image_resize = image.resize(size, Image.LANCZOS)
256
  temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='1.png')
257
  temp_filename = temp_file.name
258
  image_resize.save(temp_filename, dpi=(300, 300))