Text_Recognition_with_Transformer_Models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8f99dd35-22ea-4be8-b231-86d952935396",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fri Mar 29 19:12:35 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.161.07             Driver Version: 535.161.07   CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA GeForce RTX 4080        Off | 00000000:2D:00.0 Off |                  N/A |\n",
      "|  0%   37C    P3              24W / 320W |     89MiB / 16376MiB |      0%      Default |\n",
      "|                                         |                      |                  N/A |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|    0   N/A  N/A      4415      G   /usr/lib/xorg/Xorg                           81MiB |\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec16485c-de3a-4e8b-bfc6-b58425285847",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Could not find image processor class in the image processor config or the model config. Loading based on pattern matching with the model's feature extractor configuration. Please open a PR/issue to update `preprocessor_config.json` to use `image_processor_type` instead of `feature_extractor_type`. This warning will be removed in v4.40.\n",
      "Some weights of VisionEncoderDecoderModel were not initialized from the model checkpoint at microsoft/trocr-base-handwritten and are newly initialized: ['encoder.pooler.dense.bias', 'encoder.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from PIL import Image\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from transformers import TrOCRProcessor, VisionEncoderDecoderModel, TrainingArguments, Trainer\n",
    "from torch.utils.data.dataset import random_split\n",
    "class SpanishDocumentsDataset(Dataset):\n",
    "    def __init__(self, image_dir, text_dir, processor):\n",
    "        self.image_dir = image_dir\n",
    "        self.text_dir = text_dir\n",
    "        self.processor = processor\n",
    "        self.filenames = [f for f in os.listdir(image_dir) if f.endswith('.jpg')]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.filenames)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        image_file = self.filenames[idx]\n",
    "        # Derive the corresponding text file name by changing the extension\n",
    "        text_file = image_file.replace('.jpg', '.txt')\n",
    "        \n",
    "        image_path = os.path.join(self.image_dir, image_file)\n",
    "        text_path = os.path.join(self.text_dir, text_file)\n",
    "        \n",
    "        image = Image.open(image_path).convert(\"RGB\")\n",
    "        pixel_values = self.processor(image, return_tensors=\"pt\").pixel_values.squeeze()\n",
    "\n",
    "        with open(text_path, 'r', encoding='utf-8') as file:\n",
    "            text = file.read().strip()\n",
    "        labels = self.processor.tokenizer(text, return_tensors=\"pt\").input_ids.squeeze()\n",
    "        \n",
    "        # Treat padding specially for label calculation, if necessary\n",
    "        labels[labels == self.processor.tokenizer.pad_token_id] = -100\n",
    "        \n",
    "        return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "\n",
    "\n",
    "image_dir = 'line_data/line_images'\n",
    "text_dir = 'line_data/line_texts'\n",
    "\n",
    "def collate_fn(batch):\n",
    "    pixel_values = [item['pixel_values'] for item in batch]\n",
    "    labels = [item['labels'] for item in batch]\n",
    "    \n",
    "    # Padding value for labels should be -100 to ignore tokens during loss calculation\n",
    "    labels = pad_sequence(labels, batch_first=True, padding_value=-100)\n",
    "    pixel_values = torch.stack(pixel_values)\n",
    "    \n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "\n",
    "\n",
    "# Initialize processor and model with correct configurations\n",
    "processor = TrOCRProcessor.from_pretrained(\"microsoft/trocr-base-handwritten\")\n",
    "model = VisionEncoderDecoderModel.from_pretrained(\"microsoft/trocr-base-handwritten\")\n",
    "model.config.decoder_start_token_id = processor.tokenizer.bos_token_id\n",
    "model.config.pad_token_id = processor.tokenizer.pad_token_id\n",
    "model.config.eos_token_id = processor.tokenizer.eos_token_id\n",
    "model.config.ignore_pad_token_for_loss = True  # Ensure pad tokens are ignored during loss calculation\n",
    "\n",
    "# Prepare dataset and data loader\n",
    "dataset = SpanishDocumentsDataset(image_dir, text_dir, processor=processor)\n",
    "\n",
    "dataset_size = len(dataset)\n",
    "eval_size = int(dataset_size * 0.1)  # 10% for evaluation\n",
    "train_size = dataset_size - eval_size  # Remaining for training\n",
    "\n",
    "# Split the dataset\n",
    "train_dataset, eval_dataset = random_split(dataset, [train_size, eval_size])\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)\n",
    "eval_loader = DataLoader(eval_dataset, batch_size=8, collate_fn=collate_fn)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "43bd1da2-83ed-4e30-bd28-a88de507ec2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 3, 384, 384]) torch.Size([8, 16])\n"
     ]
    }
   ],
   "source": [
    "# # To see how the data is loaded, let's fetch a batch from the DataLoader\n",
    "for batch in train_loader:\n",
    "    print(batch['pixel_values'].shape, batch['labels'].shape)\n",
    "    break  # Just show the first batch for demonstration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "765c2efd-432d-4277-9f6e-bf98ace72b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Levenshtein as lev\n",
    "\n",
    "def compute_cer(decoded_preds, decoded_labels):\n",
    "    \"\"\"\n",
    "    Compute the Character Error Rate (CER).\n",
    "    CER is defined as the edit distance between the predicted and true sequences\n",
    "    divided by the length of the true sequence.\n",
    "    \"\"\"\n",
    "    total_edit_distance = 0\n",
    "    total_length = 0\n",
    "    \n",
    "    for pred, label in zip(decoded_preds, decoded_labels):\n",
    "        total_edit_distance += lev.distance(pred, label)\n",
    "        total_length += len(label)\n",
    "    \n",
    "    cer = total_edit_distance / total_length if total_length > 0 else 0\n",
    "    return cer\n",
    "\n",
    "def compute_wer(decoded_preds, decoded_labels):\n",
    "    \"\"\"\n",
    "    Compute the Word Error Rate (WER).\n",
    "    WER is defined as the edit distance between the predicted and true sequences\n",
    "    of words divided by the number of words in the true sequence.\n",
    "    \"\"\"\n",
    "    total_edit_distance = 0\n",
    "    total_words = 0\n",
    "    \n",
    "    for pred, label in zip(decoded_preds, decoded_labels):\n",
    "        pred_words = pred.split()\n",
    "        label_words = label.split()\n",
    "        \n",
    "        total_edit_distance += lev.distance(pred_words, label_words)\n",
    "        total_words += len(label_words)\n",
    "    \n",
    "    wer = total_edit_distance / total_words if total_words > 0 else 0\n",
    "    return wer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9a2e4ac0-6492-4fe4-8b0e-bb62354f269f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_308040/1290989381.py:5: FutureWarning: load_metric is deprecated and will be removed in the next major version of datasets. Use 'evaluate.load' instead, from the new library 🤗 Evaluate: https://huggingface.co/docs/evaluate\n",
      "  cer_metric = load_metric(\"cer\")\n",
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/datasets/load.py:756: FutureWarning: The repository for cer contains custom code which must be executed to correctly load the metric. You can inspect the repository content at https://raw.githubusercontent.com/huggingface/datasets/2.18.0/metrics/cer/cer.py\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this metric from the next major release of `datasets`.\n",
      "  warnings.warn(\n",
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/datasets/load.py:756: FutureWarning: The repository for wer contains custom code which must be executed to correctly load the metric. You can inspect the repository content at https://raw.githubusercontent.com/huggingface/datasets/2.18.0/metrics/wer/wer.py\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this metric from the next major release of `datasets`.\n",
      "  warnings.warn(\n",
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/transformers/optimization.py:457: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/transformers/models/trocr/processing_trocr.py:136: FutureWarning: `feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.\n",
      "  warnings.warn(\n",
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/accelerate/accelerator.py:432: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
      "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1220' max='1220' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1220/1220 1:21:22, Epoch 20/20]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Cer</th>\n",
       "      <th>Wer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>1.976300</td>\n",
       "      <td>1.183152</td>\n",
       "      <td>0.210676</td>\n",
       "      <td>0.406832</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.531900</td>\n",
       "      <td>0.441966</td>\n",
       "      <td>0.080636</td>\n",
       "      <td>0.217391</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>0.332400</td>\n",
       "      <td>0.414173</td>\n",
       "      <td>0.120954</td>\n",
       "      <td>0.239130</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.341300</td>\n",
       "      <td>0.365381</td>\n",
       "      <td>0.124929</td>\n",
       "      <td>0.260870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>250</td>\n",
       "      <td>0.433400</td>\n",
       "      <td>0.502178</td>\n",
       "      <td>0.089154</td>\n",
       "      <td>0.214286</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>1.279700</td>\n",
       "      <td>2.216505</td>\n",
       "      <td>0.331630</td>\n",
       "      <td>0.649068</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>350</td>\n",
       "      <td>0.255500</td>\n",
       "      <td>0.565370</td>\n",
       "      <td>0.141965</td>\n",
       "      <td>0.322981</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>0.307900</td>\n",
       "      <td>0.535151</td>\n",
       "      <td>0.097672</td>\n",
       "      <td>0.257764</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>450</td>\n",
       "      <td>0.573100</td>\n",
       "      <td>0.681593</td>\n",
       "      <td>0.120386</td>\n",
       "      <td>0.316770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.782000</td>\n",
       "      <td>0.893538</td>\n",
       "      <td>0.158433</td>\n",
       "      <td>0.357143</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>550</td>\n",
       "      <td>0.615600</td>\n",
       "      <td>1.061507</td>\n",
       "      <td>0.168654</td>\n",
       "      <td>0.369565</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.569500</td>\n",
       "      <td>0.883370</td>\n",
       "      <td>0.214651</td>\n",
       "      <td>0.437888</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>650</td>\n",
       "      <td>0.302200</td>\n",
       "      <td>0.552620</td>\n",
       "      <td>0.135718</td>\n",
       "      <td>0.304348</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>0.319700</td>\n",
       "      <td>0.791597</td>\n",
       "      <td>0.151618</td>\n",
       "      <td>0.344720</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>750</td>\n",
       "      <td>0.184700</td>\n",
       "      <td>0.702757</td>\n",
       "      <td>0.175468</td>\n",
       "      <td>0.369565</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.210200</td>\n",
       "      <td>0.512535</td>\n",
       "      <td>0.118115</td>\n",
       "      <td>0.270186</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>850</td>\n",
       "      <td>0.120800</td>\n",
       "      <td>0.560051</td>\n",
       "      <td>0.104486</td>\n",
       "      <td>0.270186</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>0.104700</td>\n",
       "      <td>0.614946</td>\n",
       "      <td>0.123793</td>\n",
       "      <td>0.298137</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>950</td>\n",
       "      <td>0.161600</td>\n",
       "      <td>0.683395</td>\n",
       "      <td>0.105622</td>\n",
       "      <td>0.267081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.066600</td>\n",
       "      <td>0.516289</td>\n",
       "      <td>0.110165</td>\n",
       "      <td>0.254658</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1050</td>\n",
       "      <td>0.043100</td>\n",
       "      <td>0.490517</td>\n",
       "      <td>0.090857</td>\n",
       "      <td>0.226708</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>0.021000</td>\n",
       "      <td>0.453972</td>\n",
       "      <td>0.095400</td>\n",
       "      <td>0.232919</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1150</td>\n",
       "      <td>0.021700</td>\n",
       "      <td>0.438193</td>\n",
       "      <td>0.088586</td>\n",
       "      <td>0.220497</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>0.000800</td>\n",
       "      <td>0.427375</td>\n",
       "      <td>0.088018</td>\n",
       "      <td>0.220497</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Removed shared tensor {'decoder.output_projection.weight'} while saving. This should be OK, but check by verifying that you don't receive any warning while reloading\n",
      "There were missing keys in the checkpoint model loaded: ['decoder.output_projection.weight'].\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_metric\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "from transformers import AdamW\n",
    "\n",
    "cer_metric = load_metric(\"cer\")\n",
    "wer_metric = load_metric(\"wer\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    # Extract the logits from the tuple if necessary\n",
    "    if isinstance(logits, tuple):\n",
    "        logits = logits[0]\n",
    "\n",
    "    # Convert logits to the most likely token IDs\n",
    "    predictions = logits.argmax(-1)\n",
    "    \n",
    "    # Decode predictions to text\n",
    "    decoded_preds = processor.batch_decode(predictions, skip_special_tokens=True)\n",
    "    \n",
    "    # Prepare labels for decoding\n",
    "    decoded_labels = []\n",
    "    for label in labels:\n",
    "        # Filter out -100 values which are used for padding/ignored indices\n",
    "        label_filtered = [token for token in label if token != -100]\n",
    "        decoded_label = processor.decode(label_filtered, skip_special_tokens=True)\n",
    "        decoded_labels.append(decoded_label)\n",
    "\n",
    "    cer = compute_cer(decoded_preds, decoded_labels)  \n",
    "    wer = compute_wer(decoded_preds, decoded_labels) \n",
    "\n",
    "    return {\"cer\": cer, \"wer\": wer}\n",
    "\n",
    "\n",
    "total_train_steps = (len(train_dataset) // 8) * 20\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)\n",
    "scheduler = get_linear_schedule_with_warmup(optimizer, \n",
    "                                            num_warmup_steps=500, \n",
    "                                            num_training_steps=total_train_steps)\n",
    "\n",
    "\n",
    "# TrainingArguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./trocr_finetuned\",\n",
    "    per_device_train_batch_size=8,\n",
    "    gradient_accumulation_steps=1,\n",
    "    num_train_epochs=20,\n",
    "    logging_dir='./logs',\n",
    "    logging_steps=10,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=50,\n",
    "    save_steps=50,\n",
    "    warmup_steps=500,\n",
    "    weight_decay=0.01,\n",
    "    save_total_limit=2,\n",
    "    load_best_model_at_end=True,\n",
    ")\n",
    "\n",
    "# Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,  \n",
    "    eval_dataset=eval_dataset,  \n",
    "    optimizers=(optimizer, scheduler),\n",
    "    data_collator=collate_fn,\n",
    "    tokenizer=processor.feature_extractor,\n",
    "    compute_metrics=compute_metrics, \n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train()\n",
    "\n",
    "# Save the fine-tuned model and processor\n",
    "trainer.save_model(\"./trocr_finetuned\")\n",
    "processor.save_pretrained(\"./trocr_finetuned\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cbf0004f-fc7d-4600-83c7-b2dc565cb143",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing page_26:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/vol/bitbucket/ur23/g/lib/python3.10/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  page_26_line_segment_0.jpg: Nobleza Vrvuosa.\n",
      "  page_26_line_segment_1.jpg: Si por euitar un pegado mortal\n",
      "  page_26_line_segment_2.jpg: queys de poner vuestra vida en pe-\n",
      "  page_26_line_segment_3.jpg: ligro, atriesgalda, que es el mejor\n",
      "  page_26_line_segment_4.jpg: empleo que della podeys hazer, y\n",
      "  page_26_line_segment_5.jpg: de vuestra hazienda para este sin en\n",
      "  page_26_line_segment_6.jpg: redemir cautuos, y sacar mugenes\n",
      "  page_26_line_segment_7.jpg: de pegado, dotandolas liberalmen-\n",
      "  page_26_line_segment_8.jpg: Caton dixo, nunca hagas el bien\n",
      "  page_26_line_segment_9.jpg: porque se sepa, dad pues vos sin\n",
      "  page_26_line_segment_10.jpg: bueno a cualquiera obra, con que,\n",
      "  page_26_line_segment_11.jpg: huyreys de la hijocresia, pero tam-\n",
      "  page_26_line_segment_12.jpg: poco escondays las que han de ser\n",
      "  page_26_line_segment_13.jpg: de buen exemplo, pues es obliga-\n",
      "  page_26_line_segment_14.jpg: ion de personas tales el darlo, y\n",
      "  page_26_line_segment_15.jpg: Ontario tentación en algunos. No.\n",
      "  page_26_line_segment_16.jpg: hagays profession de santero, pe-\n",
      "  page_26_line_segment_17.jpg: ro si de buen Christiano, no apro-\n",
      "  page_26_line_segment_18.jpg: veys (mas tampoco reproveys) Santi-\n",
      "  page_26_line_segment_19.jpg: dades dodosas, sino estimad las cier-\n",
      "  page_26_line_segment_20.jpg: tas, y aprobadas, y a esto toca el no\n",
      "  page_26_line_segment_21.jpg: ser milagrero. Accordaos del Rey S.\n",
      "  page_26_line_segment_22.jpg: Luys que no quiso ver con los ojos\n",
      "  page_26_line_segment_23.jpg: no que mejor veya con la Fe.\n",
      "  page_26_line_segment_24.jpg: 707.\n",
      "\n",
      "Full text for page_26:\n",
      "Nobleza Vrvuosa.\n",
      "Si por euitar un pegado mortal\n",
      "queys de poner vuestra vida en pe-\n",
      "ligro, atriesgalda, que es el mejor\n",
      "empleo que della podeys hazer, y\n",
      "de vuestra hazienda para este sin en\n",
      "redemir cautuos, y sacar mugenes\n",
      "de pegado, dotandolas liberalmen-\n",
      "Caton dixo, nunca hagas el bien\n",
      "porque se sepa, dad pues vos sin\n",
      "bueno a cualquiera obra, con que,\n",
      "huyreys de la hijocresia, pero tam-\n",
      "poco escondays las que han de ser\n",
      "de buen exemplo, pues es obliga-\n",
      "ion de personas tales el darlo, y\n",
      "Ontario tentación en algunos. No.\n",
      "hagays profession de santero, pe-\n",
      "ro si de buen Christiano, no apro-\n",
      "veys (mas tampoco reproveys) Santi-\n",
      "dades dodosas, sino estimad las cier-\n",
      "tas, y aprobadas, y a esto toca el no\n",
      "ser milagrero. Accordaos del Rey S.\n",
      "Luys que no quiso ver con los ojos\n",
      "no que mejor veya con la Fe.\n",
      "707.\n",
      "\n",
      "==================================================\n",
      "\n",
      "Processing page_27:\n",
      "  page_27_line_segment_0.jpg: Nobleza Vintuosa. 27\n",
      "  page_27_line_segment_1.jpg: imitad siempre aquella esperan\n",
      "  page_27_line_segment_2.jpg: ca en Dios (tan bien probada) del-\n",
      "  page_27_line_segment_3.jpg: gian Patriarcha Abraham, y hazed\n",
      "  page_27_line_segment_4.jpg: como dezia el Santo Padre Franco-\n",
      "  page_27_line_segment_5.jpg: co de Borja todas las possibles dili-\n",
      "  page_27_line_segment_6.jpg: gencias en los negocios, como sino\n",
      "  page_27_line_segment_7.jpg: haulviera Dios, pero no siando en nin-\n",
      "  page_27_line_segment_8.jpg: guna, sino solo en el. No comuni\n",
      "  page_27_line_segment_9.jpg: queys Astrologos, que no ay certi-\n",
      "  page_27_line_segment_10.jpg: dumbre,antes confusion, y mil tro-\n",
      "  page_27_line_segment_11.jpg: piezos en su ciencia, mas dado caso\n",
      "  page_27_line_segment_12.jpg: que fuera segura si anncian algunn\n",
      "  page_27_line_segment_13.jpg: bien, y ha de venir, es tormento el\n",
      "  page_27_line_segment_14.jpg: esperarle, y se estima en menos que\n",
      "  page_27_line_segment_15.jpg: do llega, y si no sale cierto, lo es la\n",
      "  page_27_line_segment_16.jpg: pena de tal enzano, pues si mal, para\n",
      "  page_27_line_segment_17.jpg: que se ha de anticipar el sentimien-\n",
      "  page_27_line_segment_18.jpg: to deljaciendo de llegar, y sino por-\n",
      "  page_27_line_segment_19.jpg: que ha de congojar lo que nunca\n",
      "  page_27_line_segment_20.jpg: aimad sobre todo a Dios, estando,\n",
      "  page_27_line_segment_21.jpg: dispusto a dar por su Fe, y honra la)\n",
      "  page_27_line_segment_22.jpg: vida, como muchos Reyes, y Princi-\n",
      "  page_27_line_segment_23.jpg: pes, a quien illustro incomparable-\n",
      "\n",
      "Full text for page_27:\n",
      "Nobleza Vintuosa. 27\n",
      "imitad siempre aquella esperan\n",
      "ca en Dios (tan bien probada) del-\n",
      "gian Patriarcha Abraham, y hazed\n",
      "como dezia el Santo Padre Franco-\n",
      "co de Borja todas las possibles dili-\n",
      "gencias en los negocios, como sino\n",
      "haulviera Dios, pero no siando en nin-\n",
      "guna, sino solo en el. No comuni\n",
      "queys Astrologos, que no ay certi-\n",
      "dumbre,antes confusion, y mil tro-\n",
      "piezos en su ciencia, mas dado caso\n",
      "que fuera segura si anncian algunn\n",
      "bien, y ha de venir, es tormento el\n",
      "esperarle, y se estima en menos que\n",
      "do llega, y si no sale cierto, lo es la\n",
      "pena de tal enzano, pues si mal, para\n",
      "que se ha de anticipar el sentimien-\n",
      "to deljaciendo de llegar, y sino por-\n",
      "que ha de congojar lo que nunca\n",
      "aimad sobre todo a Dios, estando,\n",
      "dispusto a dar por su Fe, y honra la)\n",
      "vida, como muchos Reyes, y Princi-\n",
      "pes, a quien illustro incomparable-\n",
      "\n",
      "==================================================\n",
      "\n",
      "Processing page_28:\n",
      "  page_28_line_segment_0.jpg: mente mas el derramar por esta cau-\n",
      "  page_28_line_segment_1.jpg: sa su sangre, que el averla heredado\n",
      "  page_28_line_segment_2.jpg: tan generosa, y vivid determinado\n",
      "  page_28_line_segment_3.jpg: de no perder ocasión en servirle, y\n",
      "  page_28_line_segment_4.jpg: Austro Confessor (que escoge-\n",
      "  page_28_line_segment_5.jpg: reys esiritual, docto, y hombre de\n",
      "  page_28_line_segment_6.jpg: gran talento) tened mucho respeto,\n",
      "  page_28_line_segment_7.jpg: ydadle autoridad para que os diga\n",
      "  page_28_line_segment_8.jpg: libremente quantas verdades a vue-\n",
      "  page_28_line_segment_9.jpg: stra alma importen, en las cosas to-\n",
      "  page_28_line_segment_10.jpg: cantes a ella, obedecedle entera-\n",
      "  page_28_line_segment_11.jpg: mente con todo rendimiento, y tal\n",
      "  page_28_line_segment_12.jpg: que no admitays razon para lo que\n",
      "  page_28_line_segment_13.jpg: los ordenare en esias, materias por\n",
      "  page_28_line_segment_14.jpg: no perder el merito de la Fe, obre-\n",
      "  page_28_line_segment_15.jpg: diencia ciega (que aquila deve aver)\n",
      "  page_28_line_segment_16.jpg: tomad su consejo, pues esco giendo-\n",
      "  page_28_line_segment_17.jpg: le con las partes dichas no aura peli-\n",
      "  page_28_line_segment_18.jpg: gro de que abuse desto, metiendose\n",
      "  page_28_line_segment_19.jpg: en el gobierno de todo, y querien-\n",
      "  page_28_line_segment_20.jpg: do consequir lo que pidiere justo, a\n",
      "  page_28_line_segment_21.jpg: injusto (que es propiedad de igno-\n",
      "  page_28_line_segment_22.jpg: rantes, y no muy espirituales) y\n",
      "\n",
      "Full text for page_28:\n",
      "mente mas el derramar por esta cau-\n",
      "sa su sangre, que el averla heredado\n",
      "tan generosa, y vivid determinado\n",
      "de no perder ocasión en servirle, y\n",
      "Austro Confessor (que escoge-\n",
      "reys esiritual, docto, y hombre de\n",
      "gran talento) tened mucho respeto,\n",
      "ydadle autoridad para que os diga\n",
      "libremente quantas verdades a vue-\n",
      "stra alma importen, en las cosas to-\n",
      "cantes a ella, obedecedle entera-\n",
      "mente con todo rendimiento, y tal\n",
      "que no admitays razon para lo que\n",
      "los ordenare en esias, materias por\n",
      "no perder el merito de la Fe, obre-\n",
      "diencia ciega (que aquila deve aver)\n",
      "tomad su consejo, pues esco giendo-\n",
      "le con las partes dichas no aura peli-\n",
      "gro de que abuse desto, metiendose\n",
      "en el gobierno de todo, y querien-\n",
      "do consequir lo que pidiere justo, a\n",
      "injusto (que es propiedad de igno-\n",
      "rantes, y no muy espirituales) y\n",
      "\n",
      "==================================================\n",
      "\n",
      "Processing page_29:\n",
      "  page_29_line_segment_0.jpg: mas de huyr desto, ganareys el dar\n",
      "  page_29_line_segment_1.jpg: Credito, y autoridad a todas vuestras\n",
      "  page_29_line_segment_2.jpg: acciones eligiendole con ella, si bio-\n",
      "  page_29_line_segment_3.jpg: la negocación de las cosas pias,\n",
      "  page_29_line_segment_4.jpg: toca principalmente al Confessor,\n",
      "  page_29_line_segment_5.jpg: elqual siendo a proposito no muda-\n",
      "  page_29_line_segment_6.jpg: reys, sino a mas no poder, y aquos\n",
      "  page_29_line_segment_7.jpg: aduiento, que aunq a todas las Reli-\n",
      "  page_29_line_segment_8.jpg: iones tenzays el amor que estadi-\n",
      "  page_29_line_segment_9.jpg: cho, en esta materia no os ateysa\n",
      "  page_29_line_segment_10.jpg: ninguma, escaped Confessor dond\n",
      "  page_29_line_segment_11.jpg: lo halleys mas conuiniente, que un\n",
      "  page_29_line_segment_12.jpg: sujeto, he ha de buscar para esto, y no\n",
      "  page_29_line_segment_13.jpg: El centesaros podria ser cada\n",
      "  page_29_line_segment_14.jpg: ocho dias, y comulgar quantdo al\n",
      "  page_29_line_segment_15.jpg: Confessor parezca, desembaracan-\n",
      "  page_29_line_segment_16.jpg: doos aquella mainana de cualquiera\n",
      "  page_29_line_segment_17.jpg: otra octupación, y no vseys de almo-\n",
      "  page_29_line_segment_18.jpg: Bien me pareceria rezasedes el\n",
      "  page_29_line_segment_19.jpg: Oficio divino, si las occupiones\n",
      "  page_29_line_segment_20.jpg: obligatorias os diessen lugar,\n",
      "  page_29_line_segment_21.jpg: lo melos el dela virgen, y su Rosa\n",
      "  page_29_line_segment_22.jpg: rio-\n",
      "\n",
      "Full text for page_29:\n",
      "mas de huyr desto, ganareys el dar\n",
      "Credito, y autoridad a todas vuestras\n",
      "acciones eligiendole con ella, si bio-\n",
      "la negocación de las cosas pias,\n",
      "toca principalmente al Confessor,\n",
      "elqual siendo a proposito no muda-\n",
      "reys, sino a mas no poder, y aquos\n",
      "aduiento, que aunq a todas las Reli-\n",
      "iones tenzays el amor que estadi-\n",
      "cho, en esta materia no os ateysa\n",
      "ninguma, escaped Confessor dond\n",
      "lo halleys mas conuiniente, que un\n",
      "sujeto, he ha de buscar para esto, y no\n",
      "El centesaros podria ser cada\n",
      "ocho dias, y comulgar quantdo al\n",
      "Confessor parezca, desembaracan-\n",
      "doos aquella mainana de cualquiera\n",
      "otra octupación, y no vseys de almo-\n",
      "Bien me pareceria rezasedes el\n",
      "Oficio divino, si las occupiones\n",
      "obligatorias os diessen lugar,\n",
      "lo melos el dela virgen, y su Rosa\n",
      "rio-\n",
      "\n",
      "==================================================\n",
      "\n",
      "Processing page_30:\n",
      "  page_30_line_segment_0.jpg: 3, nobleza vivivosa,\n",
      "  page_30_line_segment_1.jpg: nio cada dia, cuya devocion seleslu-\n",
      "  page_30_line_segment_2.jpg: cio bion a los Emperadores Henri-\n",
      "  page_30_line_segment_3.jpg: zni, cos 1ly. VII. y a otrosynuchos.\n",
      "  page_30_line_segment_4.jpg: Lunes Oficio de divutos. Los Vier-\n",
      "  page_30_line_segment_5.jpg: nes los Palmos Penicenciales, y Os-\n",
      "  page_30_line_segment_6.jpg: cio de la Cruz, y un rato de oracion,\n",
      "  page_30_line_segment_7.jpg: procurad no perderle, pensando en\n",
      "  page_30_line_segment_8.jpg: el un para que fuystes criado, y si-\n",
      "  page_30_line_segment_9.jpg: cumplis con las obligaciones de\n",
      "  page_30_line_segment_10.jpg: Christiano, y de vuestro Estado, que\n",
      "  page_30_line_segment_11.jpg: el Emperador Carlos Quinto ocu-\n",
      "  page_30_line_segment_12.jpg: pava cada dia dos horas en este\n",
      "  page_30_line_segment_13.jpg: exercicio, en medio de sus grandes\n",
      "  page_30_line_segment_14.jpg: egocios, conociendo ser el mas im-\n",
      "  page_30_line_segment_15.jpg: Cadanoche hazed cuentas con\n",
      "  page_30_line_segment_16.jpg: Dios, y examen de vuestra concien-\n",
      "  page_30_line_segment_17.jpg: cia, pues no sabeys si amanecereys\n",
      "  page_30_line_segment_18.jpg: en el otro mundo, como sucedio a\n",
      "  page_30_line_segment_19.jpg: La Quaresima, y Semana Santa,\n",
      "  page_30_line_segment_20.jpg: mostrad con particularidad, que\n",
      "  page_30_line_segment_21.jpg: soys Christiano, celebrando con\n",
      "  page_30_line_segment_22.jpg: vestido (siempre negro) y semblante\n",
      "\n",
      "Full text for page_30:\n",
      "3, nobleza vivivosa,\n",
      "nio cada dia, cuya devocion seleslu-\n",
      "cio bion a los Emperadores Henri-\n",
      "zni, cos 1ly. VII. y a otrosynuchos.\n",
      "Lunes Oficio de divutos. Los Vier-\n",
      "nes los Palmos Penicenciales, y Os-\n",
      "cio de la Cruz, y un rato de oracion,\n",
      "procurad no perderle, pensando en\n",
      "el un para que fuystes criado, y si-\n",
      "cumplis con las obligaciones de\n",
      "Christiano, y de vuestro Estado, que\n",
      "el Emperador Carlos Quinto ocu-\n",
      "pava cada dia dos horas en este\n",
      "exercicio, en medio de sus grandes\n",
      "egocios, conociendo ser el mas im-\n",
      "Cadanoche hazed cuentas con\n",
      "Dios, y examen de vuestra concien-\n",
      "cia, pues no sabeys si amanecereys\n",
      "en el otro mundo, como sucedio a\n",
      "La Quaresima, y Semana Santa,\n",
      "mostrad con particularidad, que\n",
      "soys Christiano, celebrando con\n",
      "vestido (siempre negro) y semblante\n",
      "\n",
      "==================================================\n",
      "\n",
      "Processing page_31:\n",
      "  page_31_line_segment_0.jpg: Nobleza Virtuosa.\n",
      "  page_31_line_segment_1.jpg: que proceda de vivo sentiment\n",
      "  page_31_line_segment_2.jpg: interior la Passion de Christo, y se-\n",
      "  page_31_line_segment_3.jpg: ria bien retitaros a un Convento\n",
      "  page_31_line_segment_4.jpg: aquellos ocho dias. Servid la comi-\n",
      "  page_31_line_segment_5.jpg: daelluenes Santo a doze pobres, la\n",
      "  page_31_line_segment_6.jpg: mandoles despueslos pies, besan-\n",
      "  page_31_line_segment_7.jpg: doselos, costumbre loable de todos\n",
      "  page_31_line_segment_8.jpg: OYreys Missa cada dia, sin que\n",
      "  page_31_line_segment_9.jpg: ya ovpaacion que os lo estorve,\n",
      "  page_31_line_segment_10.jpg: que son infinitas las ganacias des-\n",
      "  page_31_line_segment_11.jpg: to, como dizen San Cirilo, y San Ci,\n",
      "  page_31_line_segment_12.jpg: priano, pero sea en la Iglesia (no en\n",
      "  page_31_line_segment_13.jpg: casajaunque en ella aveys de tener\n",
      "  page_31_line_segment_14.jpg: Oratorio muy bien adornado, y de- c,\n",
      "  page_31_line_segment_15.jpg: uoto, cuydado que tocara propria-\n",
      "  page_31_line_segment_16.jpg: mente a vuestra mucer. Pero' vos\n",
      "  page_31_line_segment_17.jpg: le tendreys, de que los Capellanes\n",
      "  page_31_line_segment_18.jpg: sean virtuosos, y no hagays esperar\n",
      "  page_31_line_segment_19.jpg: al que os ha de de dezir la Missa renes-\n",
      "  page_31_line_segment_20.jpg: tido, que es grande indecencia, ni\n",
      "  page_31_line_segment_21.jpg: seays de los que reprehende S. Agu-\n",
      "  page_31_line_segment_22.jpg: tin porque buscan Missas breves,\n",
      "  page_31_line_segment_23.jpg: Cuydad mucho de escoger los Sa-\n",
      "  page_31_line_segment_24.jpg: 20cerdotes\n",
      "\n",
      "Full text for page_31:\n",
      "Nobleza Virtuosa.\n",
      "que proceda de vivo sentiment\n",
      "interior la Passion de Christo, y se-\n",
      "ria bien retitaros a un Convento\n",
      "aquellos ocho dias. Servid la comi-\n",
      "daelluenes Santo a doze pobres, la\n",
      "mandoles despueslos pies, besan-\n",
      "doselos, costumbre loable de todos\n",
      "OYreys Missa cada dia, sin que\n",
      "ya ovpaacion que os lo estorve,\n",
      "que son infinitas las ganacias des-\n",
      "to, como dizen San Cirilo, y San Ci,\n",
      "priano, pero sea en la Iglesia (no en\n",
      "casajaunque en ella aveys de tener\n",
      "Oratorio muy bien adornado, y de- c,\n",
      "uoto, cuydado que tocara propria-\n",
      "mente a vuestra mucer. Pero' vos\n",
      "le tendreys, de que los Capellanes\n",
      "sean virtuosos, y no hagays esperar\n",
      "al que os ha de de dezir la Missa renes-\n",
      "tido, que es grande indecencia, ni\n",
      "seays de los que reprehende S. Agu-\n",
      "tin porque buscan Missas breves,\n",
      "Cuydad mucho de escoger los Sa-\n",
      "20cerdotes\n",
      "\n",
      "==================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "import os\n",
    "from PIL import Image\n",
    "import torch\n",
    "from transformers import TrOCRProcessor, VisionEncoderDecoderModel\n",
    "\n",
    "model_path = \"./trocr_finetuned\"\n",
    "processor_path = \"./trocr_finetuned\"\n",
    "\n",
    "# Load the fine-tuned model and processor\n",
    "processor = TrOCRProcessor.from_pretrained(processor_path)\n",
    "model = VisionEncoderDecoderModel.from_pretrained(model_path)\n",
    "\n",
    "# Function to generate text for a single image segment\n",
    "def generate_text_from_image_segment(image_path):\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "    pixel_values = processor(images=image, return_tensors=\"pt\").pixel_values\n",
    "    generated_ids = model.generate(pixel_values)\n",
    "    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
    "    return generated_text\n",
    "\n",
    "# Base directory where the page folders are located\n",
    "base_dir = \"test_line_segments\" \n",
    "\n",
    "\n",
    "def sort_key(filename):\n",
    "    \"\"\"\n",
    "    Custom sort function to extract the segment number from the filename\n",
    "    and use it as the key for sorting.\n",
    "    \"\"\"\n",
    "    match = re.search(r\"line_segment_(\\d+)\\.jpg\", filename)\n",
    "    if match:\n",
    "        return int(match.group(1))\n",
    "    return -1  # Return -1 if the pattern doesn't match\n",
    "\n",
    "\n",
    "# Iterate through each page's folder\n",
    "for page_folder in sorted(os.listdir(base_dir)):\n",
    "    page_path = os.path.join(base_dir, page_folder)\n",
    "    if os.path.isdir(page_path):\n",
    "        print(f\"Processing {page_folder}:\")\n",
    "        page_texts = []\n",
    "\n",
    "        # Sort the line segment images numerically based on the segment number\n",
    "        line_segment_images = sorted(os.listdir(page_path), key=sort_key)\n",
    "\n",
    "        # Iterate through each sorted line segment in the page folder\n",
    "        for line_segment_image in line_segment_images:\n",
    "            if line_segment_image.endswith('.jpg'):\n",
    "                line_segment_path = os.path.join(page_path, line_segment_image)\n",
    "                line_text = generate_text_from_image_segment(line_segment_path)\n",
    "                page_texts.append(line_text)\n",
    "                print(f\"  {line_segment_image}: {line_text}\")\n",
    "\n",
    "        # Compile and display the full page's text\n",
    "        full_page_text = \"\\n\".join(page_texts)\n",
    "        print(f\"\\nFull text for {page_folder}:\")\n",
    "        print(full_page_text)\n",
    "        print(\"\\n\" + \"=\"*50 + \"\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e6a3c1c-9854-4955-ad13-5ac59581dd21",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}