Update README.md

README.md

@@ -68,7 +68,58 @@ A CRNN+CTC version of this model trained on Kraken 4.0 (https://github.com/mitta
 Torres Aguilar, S. (2024). TRIDIS v2 : HTR model for Multilingual Medieval and Early Modern Documentary Manuscripts (11th-16th) (Version 2). Zenodo. https://doi.org/10.5281/zenodo.13862096
 
 
-The following snippet can be used to get model inferences on manuscript lines. Ideally the test dataset must be passed to the model on the form of a json file redirecting to the images:
+The following snippets can be used to get model inferences on manuscript lines. 
+
+1. Clone the model using: git lfs clone https://huggingface.co/magistermilitum/tridis_v2_HTR_historical_manuscripts
+
+2. Here is how to test the model on one single image:
+
+```python 
+from transformers import TrOCRProcessor, AutoTokenizer, VisionEncoderDecoderModel
+from safetensors.torch import load_file
+import torch.nn as nn
+
+from PIL import Image
+
+# load image from the IAM database
+path="/path/to/image/file.png"
+image = Image.open(path).convert("RGB")
+
+processor = TrOCRProcessor.from_pretrained("./tridis_v2_HTR_historical_manuscripts")
+model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-handwritten')
+
+# Load the weights of this model
+safetensors_path = "./tridis_v2_HTR_historical_manuscripts/model.safetensors" #load the weights from the downloaded model
+state_dict = load_file(safetensors_path)
+
+# Load the trocr model
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
+
+#Modify the embeddings size and vocab
+model.config.decoder.vocab_size = processor.tokenizer.vocab_size
+model.config.vocab_size = model.config.decoder.vocab_size
+model.decoder.output_projection = nn.Linear(1024, processor.tokenizer.vocab_size)
+#model.decoder.model.decoder.embed_tokens = nn.Embedding(processor.tokenizer.vocab_size, 1024, padding_idx=1)
+model.decoder.embed_tokens = nn.Embedding(processor.tokenizer.vocab_size, 1024, padding_idx=1)
+
+# set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_length = 160
+model.config.early_stopping = True
+model.config.no_repeat_ngram_size = 3
+model.config.length_penalty = 2.0
+model.config.num_beams = 3
+
+model.load_state_dict(state_dict)
+
+pixel_values = processor(images=image, return_tensors="pt").pixel_values
+
+generated_ids = model.generate(pixel_values)
+generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print(generated_text)
+```
+
+3. Here is how test the model on a dataset. Ideally the test dataset must be passed to the model on the form of a json list redirecting to the images:
 
 for ex (graphical_line_path, line_text_content):
 
@@ -78,8 +129,6 @@ for ex (graphical_line_path, line_text_content):
   etc.
 ]
 
-Clone the model using: git lfs clone https://huggingface.co/magistermilitum/tridis_v2_HTR_historical_manuscripts
-
 ```python 
 import glob
 import json, random