fix: trained model to 0.49 cer and fixed streamlit app issues

.gitignore

@@ -1,3 +1,4 @@
 __pycache__/
 core/__pycache__/
 trocr-ocr/
+light-orc/

app.py

@@ -3,17 +3,16 @@ from PIL import Image
 import torch
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 
-MODEL_NAME = 'model/'
-processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
-model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
+MODEL_PATH = 'model/'
+processor = TrOCRProcessor.from_pretrained(MODEL_PATH)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_PATH)
 
 streamlit.title('Light OCR')
 
 uploaded_file = streamlit.file_uploader('Choose an image...', type=['png', 'jpg', 'jpeg'])
 if uploaded_file:
-    image = Image.open(uploaded_file).convert("RGB")
-    image = image.resize((384, 384))
-    streamlit.image(image, caption='Uploaded Image', use_column_width=True)
+    image = Image.open(uploaded_file).convert('RGB')
+    streamlit.image(image, caption='Uploaded Image', use_container_width=True)
     
     pixel_values = processor(images=image, return_tensors='pt').pixel_values
     output_ids = model.generate(pixel_values)

config.py

@@ -2,7 +2,7 @@ DATASET_NAME = 'Teklia/IAM-line'
 MODEL_NAME = 'microsoft/trocr-small-printed'
 TRAIN_SPLIT = 'train[:200]'
 TEST_SPLIT_RATIO = 0.2
-BATCH_SIZE = 2
-EPOCHS = 1
+BATCH_SIZE = 8
+EPOCHS = 8
 OUTPUT_DIR = './trocr-ocr'
 

core/data.py

@@ -2,13 +2,15 @@ import torch
 from datasets import load_dataset
 from config import DATASET_NAME, TRAIN_SPLIT, TEST_SPLIT_RATIO
 from core.model import processor
+from PIL import Image
 
 def preprocess_batch(batch):
-    images = [img.convert("RGB") for img in batch["image"]]
-    pixel_values = processor(images=images, return_tensors="pt").pixel_values
-    batch["pixel_values"] = pixel_values
+    images = [img.convert('RGB') for img in batch["image"]]
+
+    labels = processor.tokenizer(batch['text'], padding=True, max_length=128, truncation=True).input_ids
+    pixel_values = processor.image_processor(images, return_tensors="pt").pixel_values
 
-    labels = processor.tokenizer(batch["text"], padding=True, truncation=True).input_ids
+    batch["pixel_values"] = pixel_values
     batch["labels"] = labels
 
     return batch
@@ -19,8 +21,8 @@ def load():
     train_ds = train_test['train']
     eval_ds = train_test['test']
 
-    train_ds = train_ds.map(preprocess_batch, batched = True, remove_columns = train_ds.column_names)
-    eval_ds = eval_ds.map(preprocess_batch, batched = True, remove_columns = eval_ds.column_names)
+    train_ds = train_ds.map(preprocess_batch, batched=True, remove_columns=train_ds.column_names)
+    eval_ds = eval_ds.map(preprocess_batch, batched=True, remove_columns=eval_ds.column_names)
 
     return train_ds, eval_ds
 

core/model.py

@@ -3,7 +3,7 @@ from transformers import VisionEncoderDecoderModel, TrOCRProcessor
 from config import MODEL_NAME
 
 device = torch.device('mps') if torch.backends.mps.is_available() else torch.device('cpu')
-processor = TrOCRProcessor.from_pretrained(MODEL_NAME, use_fast = False)
+processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
 model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME).to(device)
 
 model.config.decoder_start_token_id = processor.tokenizer.bos_token_id

core/train.py

@@ -1,4 +1,4 @@
-from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
 from core.model import model, processor
 from core.data import load
 from core.utils import compute_metrics
@@ -12,9 +12,11 @@ training_args = Seq2SeqTrainingArguments(
     per_device_eval_batch_size = BATCH_SIZE,
     predict_with_generate = True,
     eval_strategy = 'epoch',
-    logging_steps = 10,
+    logging_steps = 50,
     num_train_epochs = EPOCHS,
     save_total_limit = 1,
+    remove_unused_columns = False,
+    learning_rate = 5e-5,
     fp16 = False
 )
 
@@ -24,9 +26,14 @@ trainer = Seq2SeqTrainer(
     train_dataset = train_ds,
     eval_dataset = eval_ds,
     processing_class = processor.image_processor,
-    data_collator = default_data_collator,
     compute_metrics = compute_metrics
 )
 
-if __name__ == '__main__':
+def train_save():
     trainer.train()
+    trainer.save_model('./model')
+    processor.save_pretrained('./model')
+
+
+if __name__ == '__main__':
+    train_save()

model/generation_config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:66a55c9bff6d80e77c8deb6dba8dd79d867da689c7e0f1e1eddb265f8a92fb1b
-size 185
+oid sha256:91fde0da8b70ba657bd5e495956d6661ebf5ed65daeb70a4bcb488e9c62c046a
+size 155

model/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6d77161e6a5564a2d70e53b5dabfad12b67fb2e9bd7c3cc7555b1fe056bc8826
+oid sha256:2d7f1b418d01098f6a8c33b290b12952b49235e97d666feea5a95b5e38c250a3
 size 246430696

model/preprocessor_config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5172006ffcaf0f407db91ac4ada30ad6ca86183fa37ce6059f966bfaffb880cb
-size 411
+oid sha256:36a945a7cc645688b9ef64dabae16979cf5f7c1c448569cc306694edc0598b9b
+size 450

model/sentencepiece.bpe.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f5e2fefcf793761a76a6bfb8ad35489f9c203b25557673284b6d032f41043f4
+size 1356293

model/special_tokens_map.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5469a60db23249c7f8945013d78df30b44b6bf686c6bb4740f4223f77b1b535
+size 279

model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:969a92d9be8996720f3523976fe57f101a56d920b388707d48641055596c114f
+size 4494958

model/tokenizer_config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0836517b3d82dcc162f06172e0b50bc1df3024cce7cf2d71ed009acd4d8c75ea
+size 1268

model/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26b9bcd85dc494ff7ae55bc22488e87174ffd31d5b8205496c292075416de299
+size 5432

requirements.txt

@@ -5,5 +5,6 @@ datasets
 evaluate
 jiwer
 Pillow
-accelerator
+accelerate
 streamlit
+sentencepiece