Views
No views yet
pip install torch transformers pillow1from PIL import Image
2
3def preprocess_image(image):
4 target_size = (224, 224)
5 original_size = image.size
6
7 aspect_ratio = original_size[0] / original_size[1]
8 if aspect_ratio > 1:
9 new_width = target_size[0]
10 new_height = int(target_size[0] / aspect_ratio)
11 else:
12 new_height = target_size[1]
13 new_width = int(target_size[1] * aspect_ratio)
14
15 resized_img = image.resize((new_width, new_height))
16
17 padding_width = target_size[0] - new_width
18 padding_height = target_size[1] - new_height
19 pad_left = padding_width // 2
20 pad_top = padding_height // 2
21
22 pad_image = Image.new('RGB', target_size, (255, 255, 255))
23 pad_image.paste(resized_img, (pad_left, pad_top))
24 return pad_image1import torch
2from PIL import Image
3from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTFeatureExtractor, TrOCRProcessor
4
5# Load the model and processor
6tokenizer = AutoTokenizer.from_pretrained("aayushpuri01/TrOCR-Devanagari")
7model1 = VisionEncoderDecoderModel.from_pretrained("aayushpuri01/TrOCR-Devanagari")
8feature_extractor1 = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
9processor1 = TrOCRProcessor(feature_extractor=feature_extractor1, tokenizer=tokenizer)
10
11device = 'cuda' if torch.cuda.is_available() else 'cpu'
12model1.to(device)
13
14# Prediction function
15def predict(image):
16 # Preprocess the image
17 image = Image.open(image).convert("RGB")
18 image = preprocess_image(image)
19 pixel_values = processor1(image, return_tensors="pt").pixel_values.to(device)
20
21 # Generate text from the image
22 generated_ids = model1.generate(pixel_values)
23 generated_text = processor1.batch_decode(generated_ids, skip_special_tokens=True)[0]
24
25 return generated_text1# Load and predict
2image_path = "path_to_your_image.jpg"
3predicted_text = predict(image_path)
4print("Predicted Text:", predicted_text)1training_args = Seq2SeqTrainingArguments(
2 predict_with_generate=True,
3 evaluation_strategy="steps",
4 per_device_train_batch_size=32,
5 per_device_eval_batch_size=32,
6 output_dir='/workspace/checkpoint-save/',
7 save_total_limit=2,
8 logging_steps=2,
9 save_steps=1000,
10 eval_steps=1000,
11 save_strategy="steps",
12 load_best_model_at_end=True,
13 metric_for_best_model="cer",
14 greater_is_better=False,
15 num_train_epochs=15
16)