Upload 3 files

augments.py

@@ -0,0 +1,42 @@
+import torch
+import random
+import torchvision.transforms as T
+import numpy as np
+
+class RandAug:
+    """Randomly chosen image augmentations."""
+
+    def __init__(self):
+        # Augmentation options
+        self.trans = ['rotation', 'blur', 'color', 'sharpness']
+        
+    def __call__(self, img):
+        # Randomly choose the number of augmentations used for input image 
+        n_transforms = random.randint(1, len(self.trans))
+        # Randomly choose the augmentation types
+        transforms = random.sample(self.trans, n_transforms)
+        
+        # Implement the augmentations sequentially
+        if 'rotation' in transforms:
+            rotation = random.randint(-10, 10)
+            img = T.functional.rotate(img=img, angle=rotation, expand=True, fill=255)         
+        
+        if 'blur' in transforms:
+            kernel = random.choice([1,3,5])
+            transform = T.GaussianBlur(kernel, sigma=(0.1, 2.0))  
+            img = transform(img)
+
+        if 'color' in transforms:
+            rand_brightness = random.uniform(0, 0.3)
+            rand_hue = random.uniform(0, 0.5)
+            rand_contrast = random.uniform(0, 0.5)
+            rand_saturation = random.uniform(0, 0.5)
+            transform = T.ColorJitter(brightness=rand_brightness, contrast=rand_contrast, saturation=rand_saturation, hue=rand_hue)
+            img = transform(img)
+
+        if 'sharpness' in transforms:
+            sharpness = 1+(np.random.exponential()/2)
+            trans = T.RandomAdjustSharpness(sharpness, p=1)
+            img = transform(img)
+
+        return img 

dataset.py

@@ -0,0 +1,44 @@
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+from torch.utils.data import Dataset
+
+from augments import RandAug
+
+# Torch dataset
+class TextlineDataset(Dataset):
+    def __init__(self, root_dir, df, processor, augment=False, max_target_length=128):
+        self.root_dir = root_dir
+        self.df = df
+        self.processor = processor
+        self.augment = augment
+        self.augmentator = RandAug()
+        self.max_target_length = max_target_length
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        # get file name + text 
+        file_name = self.df['file_name'][idx]
+        text = self.df['text'][idx]
+
+        # prepare image (i.e. resize + normalize)
+        image = Image.open(self.root_dir + file_name).convert("RGB")
+
+        # Add image augmentations
+        if self.augment:
+            image = self.augmentator(image)
+
+        # extract the pixel values
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+        
+        # add labels (input_ids) by encoding the text
+        labels = self.processor.tokenizer(str(text), 
+                                          padding="max_length", truncation=True, 
+                                          max_length=self.max_target_length).input_ids
+        # important: make sure that PAD tokens are ignored by the loss function
+        labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
+        encoding = {"pixel_values": pixel_values.squeeze(), "labels": torch.tensor(labels)}
+        return encoding

train_trocr.py

@@ -0,0 +1,142 @@
+import pandas as pd
+import torch
+from PIL import Image
+import argparse
+from evaluate import load
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator, AdamW
+import torchvision.transforms as transforms
+#import torch_optimizer as optim
+
+from dataset import TextlineDataset
+
+parser = argparse.ArgumentParser('arguments for the code')
+
+parser.add_argument('--root_path', type=str, default="",
+                    help='Root path to data files.')
+parser.add_argument('--tr_data_path', type=str, default="/data/htr/trocr_data/trocr_tuomiokirjat/train/trocr/data.csv",
+                    help='Path to .csv file containing the training data.')
+parser.add_argument('--val_data_path', type=str, default="/data/htr/trocr_data/trocr_tuomiokirjat/val/trocr/data.csv",
+                    help='Path to .csv file containing the validation data.')
+parser.add_argument('--output_path', type=str, default="/koodit/htr/text_recognition/trocr/tuomiokirjat/models/22112023/",
+                    help='Path for saving training results.')
+parser.add_argument('--resume_path', type=str, default="/koodit/htr/text_recognition/trocr/tuomiokirjat/models/22112023",
+                    help='Path to the previous model')
+parser.add_argument('--batch_size', type=int, default=24,
+                    help='Batch size per device.')
+parser.add_argument('--epochs', type=int, default=13,
+                    help='Number of training epochs.')
+
+args = parser.parse_args()
+
+# nohup python train_trocr.py > logs/tuomiokirjat_resume_23112023.txt 2>&1 &
+# echo $! > logs/save_pid.txt
+
+# run using 2 GPUs: torchrun --nproc_per_node=2 train_trocr.py > logs/tuomiokirjat_22112023.txt 2>&1 &
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print('Device: ', device)
+
+# Initialize processor and model
+#processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+#model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") 
+processor =TrOCRProcessor.from_pretrained(args.resume_path + "/processor")
+model = VisionEncoderDecoderModel.from_pretrained(args.resume_path + "/checkpoint-13094")
+
+model.to(device)
+
+# Initialize metrics
+cer_metric = load("cer")
+wer_metric = load("wer")
+
+# Load train and validation data to dataframes
+train_df = pd.read_csv(args.tr_data_path)
+val_df = pd.read_csv(args.val_data_path)
+#train_df = train_df.iloc[:50]
+#val_df = val_df.iloc[:10]
+
+# Reset the indices to start from zero
+train_df.reset_index(drop=True, inplace=True)
+val_df.reset_index(drop=True, inplace=True)
+
+# Create train and validation datasets
+train_dataset = TextlineDataset(root_dir=args.root_path,
+                                df=train_df,
+                                processor=processor,
+                                augment=False)
+
+eval_dataset = TextlineDataset(root_dir=args.root_path,
+                               df=val_df,
+                               processor=processor,
+                               augment=False)
+
+print("Number of training examples:", len(train_dataset))
+print("Number of validation examples:", len(eval_dataset))
+
+# Define model configuration
+
+# set special tokens used for creating the decoder_input_ids from the labels
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+# make sure vocab size is set correctly
+model.config.vocab_size = model.config.decoder.vocab_size
+# set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_length = 64
+model.config.early_stopping = True
+model.config.no_repeat_ngram_size = 3
+model.config.length_penalty = 2.0
+model.config.num_beams = 4
+
+# Set arguments for model training
+# For all argumenst see https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
+training_args = Seq2SeqTrainingArguments(
+    predict_with_generate=True,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    logging_strategy="steps",
+    logging_steps=50,
+    per_device_train_batch_size=args.batch_size,
+    per_device_eval_batch_size=args.batch_size,
+    load_best_model_at_end=True,
+    metric_for_best_model='cer',
+    greater_is_better=False,
+    #fp16=True, 
+    num_train_epochs=args.epochs,
+    save_total_limit=2,
+    output_dir=args.output_path,
+    optim="adamw_torch"
+)
+
+# Function for computing CER and WER metrics for the prediction results
+def compute_metrics(pred):
+    labels_ids = pred.label_ids
+    pred_ids = pred.predictions
+
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
+    label_str = processor.batch_decode(labels_ids, skip_special_tokens=True)
+
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+    wer = wer_metric.compute(predictions=pred_str, references=label_str)
+
+    return {"cer": cer, "wer": wer}
+
+
+# instantiate trainer
+# For all parameters see: https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainer
+trainer = Seq2SeqTrainer(
+    model=model,
+    tokenizer=processor.feature_extractor,
+    args=training_args,
+    compute_metrics=compute_metrics,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=default_data_collator,
+)
+
+# Train the model
+trainer.train()
+#trainer.train(resume_from_checkpoint = True)
+model.save_pretrained(args.output_path)
+processor.save_pretrained(args.output_path + "/processor")
+