| | --- |
| | pipeline_tag: image-to-text |
| | license: apache-2.0 |
| | --- |
| | ## Usage method: |
| | ```python |
| | from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer |
| | import torch |
| | from PIL import Image |
| | |
| | model = VisionEncoderDecoderModel.from_pretrained("AIris-Channel/vit-gpt2-verifycode-caption") |
| | feature_extractor = ViTImageProcessor.from_pretrained("AIris-Channel/vit-gpt2-verifycode-caption") |
| | tokenizer = AutoTokenizer.from_pretrained("AIris-Channel/vit-gpt2-verifycode-caption") |
| | |
| | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| | model.to(device) |
| | |
| | max_length = 16 |
| | num_beams = 4 |
| | gen_kwargs = {"max_length": max_length, "num_beams": num_beams} |
| | def predict_step(image_paths): |
| | images = [] |
| | for image_path in image_paths: |
| | i_image = Image.open(image_path) |
| | if i_image.mode != "RGB": |
| | i_image = i_image.convert(mode="RGB") |
| | |
| | images.append(i_image) |
| | |
| | pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values |
| | pixel_values = pixel_values.to(device) |
| | |
| | output_ids = model.generate(pixel_values, **gen_kwargs) |
| | |
| | preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True) |
| | preds = [pred.strip() for pred in preds] |
| | return preds |
| | |
| | pred=predict_step(['ZZZTVESE.jpg']) |
| | print(pred) #zzztvese |
| | ``` |