File size: 2,141 Bytes
b781107 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
from constants import *
from transformers import AutoTokenizer
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
def get_tokenizer():
tokenizer = AutoTokenizer.from_pretrained("gpt2")
point_tokens = [f"coord_bin_{i}" for i in range(0, NUM_BINS)]
new_tokens = [
"<point_start>", "<point_end>", "<result_start>",
"<result_end>", "<pointx_start>", "<pointx_end>",
"<pointy_start>", "<pointy_end>",
*point_tokens
]
tokenizer.add_tokens(new_tokens)
# Ensure pad token is set (GPT2 usually doesn't have one by default)
if tokenizer.pad_token is None:
tokenizer.add_special_tokens({'pad_token': '[PAD]'}) # Or use eos_token if preferred
# tokenizer.pad_token_id = tokenizer.eos_token_id # Alternative if we want padding to be EOS
print(f"Tokenizer pad token: {tokenizer.pad_token}, ID: {tokenizer.pad_token_id}")
print(f"Tokenizer EOS token: {tokenizer.eos_token}, ID: {tokenizer.eos_token_id}")
# Check if pad token ID is valid
if tokenizer.pad_token_id is None:
raise ValueError("Tokenizer pad token ID is not set!")
return tokenizer, len(tokenizer)
def image_to_tensor(image, image_size=IMAGE_SIZE):
if image.mode != 'RGB':
image = image.convert('RGB')
# We avoid the hassle of calculating
# changed co-ordinates for rotation etc for now. Can be added later.
transform = transforms.Compose([
transforms.Resize((image_size, image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGE_MEAN, std=IMAGE_STD)
])
return transform(image)
def tensor_to_image(tensor):
tensor = tensor.clone().detach()
if tensor.is_cuda:
tensor = tensor.cpu()
mean = torch.tensor(IMAGE_MEAN).view(3, 1, 1)
std = torch.tensor(IMAGE_STD).view(3, 1, 1)
tensor = tensor * std + mean
tensor = torch.clamp(tensor, 0, 1)
image_np = tensor.numpy().transpose(1, 2, 0)
image_np = (image_np * 255).astype(np.uint8)
return Image.fromarray(image_np)
tokenizer, vocab_size = get_tokenizer() # Initialize tokenizer globally |