File size: 4,042 Bytes
012f8b5 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import torch
import torch.nn as nn
from torchvision import transforms
from transformers import ViTModel, BertTokenizerFast, BertConfig, BertLMHeadModel
from PIL import Image
import os
# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Define constants
VIT_MODEL_NAME = "google/vit-base-patch16-224"
BERT_MODEL_NAME = "dbmdz/bert-base-turkish-cased"
MAX_LENGTH = 128
class ImageCaptioningModel(nn.Module):
def __init__(self, vit_model, bert_model):
super(ImageCaptioningModel, self).__init__()
self.vit = vit_model
self.bert = bert_model
self.linear = nn.Linear(self.vit.config.hidden_size, self.bert.config.hidden_size)
def forward(self, pixel_values, input_ids, attention_mask, labels=None):
image_features = self.vit(pixel_values).last_hidden_state
image_features = self.linear(image_features)
outputs = self.bert(input_ids=input_ids,
attention_mask=attention_mask,
encoder_hidden_states=image_features,
labels=labels,
return_dict=True)
return outputs.loss, outputs.logits
def load_model(model_path):
# Initialize the model components
vit_model = ViTModel.from_pretrained(VIT_MODEL_NAME)
bert_config = BertConfig.from_pretrained(BERT_MODEL_NAME)
bert_config.is_decoder = True
bert_config.add_cross_attention = True
bert_model = BertLMHeadModel.from_pretrained(BERT_MODEL_NAME, config=bert_config)
# Create the combined model
model = ImageCaptioningModel(vit_model, bert_model)
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)
model.eval()
return model
def generate_caption(model, image_path, tokenizer):
# Prepare the image
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
image = Image.open(image_path).convert('RGB')
image = transform(image).unsqueeze(0).to(device)
# Generate the caption
with torch.no_grad():
input_ids = torch.tensor([[tokenizer.cls_token_id]]).to(device)
attention_mask = torch.tensor([[1]]).to(device)
for _ in range(MAX_LENGTH):
_, logits = model(image, input_ids, attention_mask)
next_token = logits[:, -1, :].argmax(dim=-1)
if next_token.item() == tokenizer.sep_token_id:
break
input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
attention_mask = torch.cat([attention_mask, torch.tensor([[1]]).to(device)], dim=1)
caption = tokenizer.decode(input_ids[0], skip_special_tokens=True)
return caption
def main():
model_path = "./models/TeLVE_v1.1.pth"
tokenizer_path = "./tokenizer"
# Check if the model and tokenizer exist
if not os.path.exists(model_path) or not os.path.exists(tokenizer_path):
print("Model or tokenizer not found. Please make sure you have trained the model and saved it correctly.")
return
# Load the model and tokenizer
model = load_model(model_path)
tokenizer = BertTokenizerFast.from_pretrained(tokenizer_path)
# Generate captions for images in a specified directory
image_dir = "./images" # Change this to the directory containing your test images
for image_file in os.listdir(image_dir):
if image_file.lower().endswith(('.png', '.jpg', '.jpeg')):
image_path = os.path.join(image_dir, image_file)
caption = generate_caption(model, image_path, tokenizer)
print(f"Image: {image_file}")
print(f"Generated Caption: {caption}")
print("---")
if __name__ == "__main__":
main() |