Update ocr_client.py

magic_pdf/dict2md/ocr_client.py

+import configparser
 import os
 import json
 import requests
@@ -10,8 +11,8 @@ from PIL import Image
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '--url',
-        default='http://127.0.0.1:6020',
+        '--config_path',
+        default='/home/practice/magic_pdf-main/magic_pdf/config.ini',
     )
     parser.add_argument(
         '--image_path',
@@ -126,6 +127,21 @@ class PredictClient:
     def __init__(self, api_url):
         self.api_url = api_url
 
+    def check_health(self):
+        health_check_url = f'{self.api_url}/health'
+        try:
+            response = requests.get(health_check_url)
+            if response.status_code == 200:
+                logger.info("Server is healthy and ready to process requests.")
+                return True
+            else:
+                logger.error(f'Server health check failed with status code:{response.status_code}')
+                return False
+        except requests.exceptions.RequestException as e:
+            logger.error(f'Health check request failed:{e}')
+            return False
+
+
     def predict(self, image_path: str, text: str):
         payload = {
             "image_path": image_path,
@@ -143,8 +159,11 @@ class PredictClient:
 
 def main():
     args = parse_args()
-    client = PredictClient(args.url)
 
+    config = configparser.ConfigParser()
+    config.read(args.config_path)
+    ocr_server = config.get('server', 'ocr_server')
+    client = PredictClient(ocr_server)
     try:
         start_time = time.time()  # 记录开始时间
         # 压缩图片