Upload app.py

app.py

@@ -12,7 +12,6 @@ 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)
@@ -34,7 +33,7 @@ class ModelConfig:
 # 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.load_state_dict(torch.load('ocr_model_large_2024-07-25_15_32.pt', map_location=torch.device('cpu')))
 trained_model.to(device)
 trained_model.eval()
 
@@ -56,10 +55,9 @@ def ocr(image):
 # Créer l'interface Gradio
 iface = gr.Interface(
     fn=ocr,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=gr.Image(type="pil", label="Upload Image",height=300),
     outputs=gr.Textbox(label="Extracted Text"),
     title="OCR Text Extraction",
     description="Upload an image to extract text using TrOCR model."
 )
-
-iface.launch(share=True)
+iface.launch()