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