Upload 4 files

app.py

@@ -23,28 +23,43 @@ if 'models_loaded' not in st.session_state:
 @st.cache_resource
 def load_models():
     try:
-        # Load YOLO model first
-        yolo_model = YOLO('best.pt')
+        # Check if YOLO weights exist
+        if not os.path.exists('best.pt'):
+            st.error("YOLO model weights (best.pt) not found in the current directory!")
+            return None, None, None
+
+        # Load YOLO model with error handling
+        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
-        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')
+        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)}")
+            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:")

requirements.txt

@@ -5,7 +5,7 @@ Pillow==10.1.0
 transformers==4.35.2
 torch==2.1.1
 torchvision==0.16.1
-ultralytics==8.0.227
+ultralytics==8.0.196
 python-Levenshtein==0.23.0
 sentencepiece==0.1.99
 protobuf==4.25.1