Upload 4 files

app.py

@@ -3,6 +3,7 @@ import os
 import numpy as np
 import cv2
 from PIL import Image
+import torch
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from ultralytics import YOLO
 import Levenshtein
@@ -21,10 +22,35 @@ if 'models_loaded' not in st.session_state:
 # Load models
 @st.cache_resource
 def load_models():
-    processor = TrOCRProcessor.from_pretrained('openthaigpt/thai-trocr')
-    ocr_model = VisionEncoderDecoderModel.from_pretrained('openthaigpt/thai-trocr')
-    yolo_model = YOLO('best.pt')  # Make sure to include this in the repository
-    return processor, ocr_model, yolo_model
+    try:
+        # Load YOLO model first
+        yolo_model = YOLO('best.pt')
+        
+        # Load TrOCR with specific model configuration
+        processor = TrOCRProcessor.from_pretrained(
+            'openthaigpt/thai-trocr',
+            use_auth_token=False,
+            trust_remote_code=True
+        )
+        
+        # Load OCR model with specific configuration
+        ocr_model = VisionEncoderDecoderModel.from_pretrained(
+            'openthaigpt/thai-trocr',
+            use_auth_token=False,
+            trust_remote_code=True
+        )
+        
+        # Move model to CPU if no GPU available
+        if not torch.cuda.is_available():
+            ocr_model = ocr_model.to('cpu')
+        
+        return processor, ocr_model, yolo_model
+    except Exception as e:
+        st.error(f"Error loading models: {str(e)}")
+        st.error("Detailed error information for debugging:")
+        import traceback
+        st.code(traceback.format_exc())
+        return None, None, None
 
 # Thai provinces list
 thai_provinces = [

requirements.txt

@@ -2,7 +2,10 @@ streamlit==1.29.0
 opencv-python-headless==4.8.1.78
 numpy==1.26.2
 Pillow==10.1.0
-transformers==4.36.2
-torch==2.1.2
+transformers==4.35.2
+torch==2.1.1
+torchvision==0.16.1
 ultralytics==8.0.227
-python-Levenshtein==0.23.0 
+python-Levenshtein==0.23.0
+sentencepiece==0.1.99
+protobuf==4.25.1