Update app.py

app.py

@@ -1,65 +1,65 @@
-import os
-import pandas as pd
-import torch
-import datetime
-import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader
-from torch import optim
-from torch.optim import AdamW
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
-from sklearn.model_selection import train_test_split
-from dataclasses import dataclass
-from PIL import Image
-from torchvision import transforms
-import matplotlib.pyplot as plt
-#from datasets import load_metric
-from tqdm.notebook import tqdm
-block_plot = False
-plt.rcParams['figure.figsize'] = (12, 9)
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
-
-model_type="large" #small|base|large
-
-@dataclass(frozen=True)
-class TrainingConfig:
-  BATCH_SIZE:    int = 15
-  EPOCHS:        int = 20
-  LEARNING_RATE: float = 0.00002
-
-@dataclass(frozen=True)
-class ModelConfig:
-  MODEL_NAME: str = 'microsoft/trocr-'+model_type+'-printed'
-
-
-# Charger le modèle entraîné à partir du fichier .pt
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-trained_model = VisionEncoderDecoderModel.from_pretrained(ModelConfig.MODEL_NAME)
-trained_model.load_state_dict(torch.load('/ocr_model_large_2024-07-25_15_32.pt'))
-trained_model.to(device)
-trained_model.eval()
-
-import gradio as gr
-import torch
-from PIL import Image
-from transformers import VisionEncoderDecoderModel, TrOCRProcessor
-import numpy as np
-
-processor = TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)
-# Fonction d'inférence
-def ocr(image):
-    image = Image.fromarray(np.array(image))  # Assurez-vous que l'image est au format PIL
-    pixel_values = processor(image, return_tensors='pt').pixel_values.to(device)
-    generated_ids = trained_model.generate(pixel_values)
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return generated_text
-
-# Créer l'interface Gradio
-iface = gr.Interface(
-    fn=ocr,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Textbox(label="Extracted Text"),
-    title="OCR Text Extraction",
-    description="Upload an image to extract text using TrOCR model."
-)
-
-iface.launch(share=True)
+import os
+import pandas as pd
+import torch
+import datetime
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torch import optim
+from torch.optim import AdamW
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
+from sklearn.model_selection import train_test_split
+from dataclasses import dataclass
+from PIL import Image
+from torchvision import transforms
+import matplotlib.pyplot as plt
+#from datasets import load_metric
+from tqdm.notebook import tqdm
+block_plot = False
+plt.rcParams['figure.figsize'] = (12, 9)
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+model_type="large" #small|base|large
+
+@dataclass(frozen=True)
+class TrainingConfig:
+  BATCH_SIZE:    int = 15
+  EPOCHS:        int = 20
+  LEARNING_RATE: float = 0.00002
+
+@dataclass(frozen=True)
+class ModelConfig:
+  MODEL_NAME: str = 'microsoft/trocr-'+model_type+'-printed'
+
+
+# Charger le modèle entraîné à partir du fichier .pt
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+trained_model = VisionEncoderDecoderModel.from_pretrained(ModelConfig.MODEL_NAME)
+trained_model.load_state_dict(torch.load('ocr_model_large_2024-07-25_15_32.pt'))
+trained_model.to(device)
+trained_model.eval()
+
+import gradio as gr
+import torch
+from PIL import Image
+from transformers import VisionEncoderDecoderModel, TrOCRProcessor
+import numpy as np
+
+processor = TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)
+# Fonction d'inférence
+def ocr(image):
+    image = Image.fromarray(np.array(image))  # Assurez-vous que l'image est au format PIL
+    pixel_values = processor(image, return_tensors='pt').pixel_values.to(device)
+    generated_ids = trained_model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
+# Créer l'interface Gradio
+iface = gr.Interface(
+    fn=ocr,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Textbox(label="Extracted Text"),
+    title="OCR Text Extraction",
+    description="Upload an image to extract text using TrOCR model."
+)
+
+iface.launch(share=True)