Upload 4 files

app.py

@@ -19,6 +19,41 @@ st.set_page_config(
 if 'models_loaded' not in st.session_state:
     st.session_state['models_loaded'] = False
 
+def load_ocr_models():
+    """Load OCR models with proper error handling"""
+    try:
+        # Set environment variables to suppress warnings
+        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+        
+        # Load processor with specific config
+        processor = TrOCRProcessor.from_pretrained(
+            'openthaigpt/thai-trocr',
+            revision='main',
+            use_auth_token=False,
+            trust_remote_code=True,
+            local_files_only=False
+        )
+        
+        # Load OCR model with specific config
+        ocr_model = VisionEncoderDecoderModel.from_pretrained(
+            'openthaigpt/thai-trocr',
+            revision='main',
+            use_auth_token=False,
+            trust_remote_code=True,
+            local_files_only=False
+        )
+        
+        # Move model to CPU explicitly
+        ocr_model = ocr_model.to('cpu')
+        
+        return processor, ocr_model
+    except Exception as e:
+        st.error(f"Error loading OCR models: {str(e)}")
+        st.error("Detailed error information:")
+        import traceback
+        st.code(traceback.format_exc())
+        return None, None
+
 # Load models
 @st.cache_resource
 def load_models():
@@ -28,41 +63,23 @@ def load_models():
             st.error("YOLO model weights (best.pt) not found in the current directory!")
             return None, None, None
 
-        # Load YOLO model with error handling
+        # Load YOLO model
         try:
             yolo_model = YOLO('best.pt', task='detect')
         except Exception as yolo_error:
             st.error(f"Error loading YOLO model: {str(yolo_error)}")
-            st.error("This might be due to a version mismatch. Please ensure you're using the correct model weights.")
             return None, None, None
         
-        # Load TrOCR with specific model configuration
-        try:
-            processor = TrOCRProcessor.from_pretrained(
-                'openthaigpt/thai-trocr',
-                use_auth_token=False,
-                trust_remote_code=True
-            )
-            
-            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')
-                
-        except Exception as ocr_error:
-            st.error(f"Error loading OCR model: {str(ocr_error)}")
+        # Load OCR models
+        processor, ocr_model = load_ocr_models()
+        if processor is None or ocr_model is None:
             return None, None, None
-        
+            
         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:")
+        st.error(f"Error in model loading: {str(e)}")
+        st.error("Detailed error information:")
         import traceback
         st.code(traceback.format_exc())
         return None, None, None

requirements.txt

@@ -2,12 +2,15 @@ streamlit==1.29.0
 opencv-python-headless==4.8.1.78
 numpy==1.26.2
 Pillow==10.1.0
-transformers==4.31.0
-tokenizers==0.13.3
-torch==2.1.1
-torchvision==0.16.1
+transformers==4.28.1
+tokenizers==0.13.2
+torch==2.0.1
+torchvision==0.15.2
 ultralytics==8.3.0
 python-Levenshtein==0.23.0
-sentencepiece==0.1.98
-protobuf==3.20.3
-huggingface-hub==0.16.4 
+sentencepiece==0.1.97
+protobuf==3.20.0
+huggingface-hub==0.14.1
+regex==2023.6.3
+requests==2.31.0
+tqdm==4.65.0