TheKOG commited on
Commit
342881d
1 Parent(s): 40053e6

Create README.md

Browse files
Files changed (1) hide show
  1. README.md +37 -0
README.md ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## Usage method:
2
+ ```python
3
+ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
4
+ import torch
5
+ from PIL import Image
6
+
7
+ model = VisionEncoderDecoderModel.from_pretrained("TheKOG/vit-gpt2-verifycode-caption")
8
+ feature_extractor = ViTImageProcessor.from_pretrained("TheKOG/vit-gpt2-verifycode-caption")
9
+ tokenizer = AutoTokenizer.from_pretrained("TheKOG/vit-gpt2-verifycode-caption")
10
+
11
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ model.to(device)
13
+
14
+ max_length = 16
15
+ num_beams = 4
16
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
17
+ def predict_step(image_paths):
18
+ images = []
19
+ for image_path in image_paths:
20
+ i_image = Image.open(image_path)
21
+ if i_image.mode != "RGB":
22
+ i_image = i_image.convert(mode="RGB")
23
+
24
+ images.append(i_image)
25
+
26
+ pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
27
+ pixel_values = pixel_values.to(device)
28
+
29
+ output_ids = model.generate(pixel_values, **gen_kwargs)
30
+
31
+ preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
32
+ preds = [pred.strip() for pred in preds]
33
+ return preds
34
+
35
+ pred=predict_step(['ZZZTVESE.jpg'])
36
+ print(pred) #zzztvese
37
+ ```